From a4221165824e5b1f5c4c10b7459f65e78dd6448d Mon Sep 17 00:00:00 2001 From: "grokkybara[bot]" <304785771+grokkybara[bot]@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:08:03 +0000 Subject: [PATCH] Synced from monorepo Synced from monorepo Changes: - Release a shell session's resources in one drop - Make the tools blocking-wait cap client-configurable and self-describing - Recognize API "exceeds budget" errors as context overflow - Retry /btw on model overload - Carry running background tasks and subagents across compaction - Require round-trip time for SDK liveness checks - Background-subagent completion reminders with a selectable delivery surface - Make a PTY shell reap itself until it reaches the registry - Recover the OS error code from a TLS-phase connection reset - Consume the attached-client signal and report why idle is withheld - Treat `.grok/sandbox.toml` edits as protected so auto mode prompts before writing - Surface history/search in the Ctrl+. cheatsheet and keep it working in history view - Delete sessions from the dashboard and welcome list - Release a session's activity record when the session ends - Stop charging auth-retry budget for fail-closed 401s; reset it across suspends - Scope skills watches on project vendor roots - Make [stop] cancel in-flight compaction - Make the leader soak measure the leader, not its harness Source-Revision: 8d69c91f02bcacf01e98d5aebbf2f92547c45738 --- Cargo.lock | 191 ++- Cargo.toml | 6 + SOURCE_REV | 2 +- clippy.toml | 1 + crates/codegen/ptyctl/src/pty.rs | 2 + crates/codegen/xai-grok-http/src/lib.rs | 132 ++ crates/codegen/xai-grok-pager-bin/Cargo.toml | 4 +- .../xai-grok-pager-pty-harness/src/pty.rs | 1 + crates/codegen/xai-grok-pager/Cargo.toml | 4 +- .../docs/user-guide/03-keyboard-shortcuts.md | 2 +- .../docs/user-guide/04-slash-commands.md | 2 +- .../docs/user-guide/17-sessions.md | 2 +- .../docs/user-guide/23-dashboard.md | 12 +- .../xai-grok-pager/src/actions/defaults.rs | 6 +- .../src/app/acp_handler/interactions.rs | 2 +- .../src/app/acp_handler/permissions.rs | 49 +- .../codegen/xai-grok-pager/src/app/actions.rs | 16 +- .../codegen/xai-grok-pager/src/app/agent.rs | 22 + .../src/app/agent_view/input.rs | 65 +- .../src/app/agent_view/interactions.rs | 6 +- .../xai-grok-pager/src/app/agent_view/mod.rs | 11 + .../src/app/agent_view/prompt.rs | 7 + .../src/app/agent_view/render.rs | 14 + .../src/app/agent_view/session.rs | 219 ++- .../xai-grok-pager/src/app/app_view.rs | 456 +++++- crates/codegen/xai-grok-pager/src/app/cli.rs | 41 + .../xai-grok-pager/src/app/dispatch/ctx.rs | 10 + .../src/app/dispatch/dashboard.rs | 292 +++- .../xai-grok-pager/src/app/dispatch/prompt.rs | 28 +- .../xai-grok-pager/src/app/dispatch/router.rs | 71 +- .../src/app/dispatch/session/foreign.rs | 49 +- .../src/app/dispatch/session/lifecycle.rs | 152 ++ .../src/app/dispatch/session/load.rs | 222 ++- .../src/app/dispatch/task_result.rs | 44 +- .../src/app/dispatch/tests/dashboard.rs | 367 ++++- .../src/app/dispatch/tests/mod.rs | 30 +- .../app/dispatch/tests/session/lifecycle.rs | 936 +++++++++++- .../src/app/dispatch/tests/session/load.rs | 13 +- .../src/app/dispatch/tests/task_result.rs | 6 +- .../xai-grok-pager/src/app/dispatch/turn.rs | 29 +- .../xai-grok-pager/src/app/effects/helpers.rs | 155 +- .../xai-grok-pager/src/app/effects/mod.rs | 22 +- .../xai-grok-pager/src/app/effects/tests.rs | 209 ++- .../xai-grok-pager/src/app/event_loop.rs | 1078 +++++++++++++- .../src/app/leader_cluster/mod.rs | 95 +- crates/codegen/xai-grok-pager/src/app/mod.rs | 54 + .../codegen/xai-grok-pager/src/app/modals.rs | 102 +- .../xai-grok-pager/src/app/session_startup.rs | 511 +++++++ .../xai-grok-pager/src/app/xt_filter.rs | 2 +- crates/codegen/xai-grok-pager/src/pty_wrap.rs | 3 +- .../xai-grok-pager/src/views/dashboard/mod.rs | 2 +- .../src/views/dashboard/render.rs | 339 +++-- .../xai-grok-pager/src/views/dashboard/row.rs | 4 +- .../src/views/dashboard/state.rs | 367 ++++- .../codegen/xai-grok-pager/src/views/modal.rs | 8 +- .../xai-grok-pager/src/views/question_view.rs | 6 + .../src/views/session_picker.rs | 83 ++ .../src/views/shortcuts_help.rs | 200 ++- .../src/views/welcome/hero_box.rs | 44 +- .../xai-grok-pager/src/views/welcome/mod.rs | 168 ++- .../src/views/welcome/workspace_mode.rs | 841 +++++++++++ .../src/actor/request_task.rs | 13 +- crates/codegen/xai-grok-sampler/src/client.rs | 155 +- crates/codegen/xai-grok-sampler/src/events.rs | 27 +- crates/codegen/xai-grok-sampler/src/handle.rs | 4 +- crates/codegen/xai-grok-sampler/src/retry.rs | 46 +- .../xai-grok-sampler/src/stream/collect.rs | 1 + .../xai-grok-sampling-types/src/error.rs | 315 +++- .../xai-grok-sampling-types/src/lib.rs | 2 +- crates/codegen/xai-grok-shell-base/Cargo.toml | 2 + crates/codegen/xai-grok-shell/CHANGELOG.md | 19 + crates/codegen/xai-grok-shell/Cargo.toml | 15 +- .../benches/skills_watcher_startup.rs | 150 ++ .../xai-grok-shell/changelogs/0.2.115.md | 4 + .../xai-grok-shell/changelogs/0.2.117.json | 37 + .../xai-grok-shell/changelogs/0.2.117.md | 18 + .../src/agent/handlers/session.rs | 4 +- .../src/agent/mvp_agent/acp_agent.rs | 63 +- .../src/agent/mvp_agent/agent_ops.rs | 721 ++++++++- .../src/agent/mvp_agent/code_nav.rs | 7 +- .../xai-grok-shell/src/agent/mvp_agent/mod.rs | 97 +- .../src/agent/mvp_agent/session_lifecycle.rs | 107 +- .../src/agent/mvp_agent/session_registry.rs | 238 +++ .../src/agent/mvp_agent/tests.rs | 379 ++++- .../src/agent/mvp_agent/tests/dhat_soak.rs | 14 +- .../xai-grok-shell/src/auth/device_code.rs | 24 +- .../codegen/xai-grok-shell/src/auth/flow.rs | 168 ++- .../xai-grok-shell/src/auth/manager.rs | 18 +- .../src/auth/manager/sleep_gate.rs | 61 +- .../xai-grok-shell/src/auth/manager_tests.rs | 6 +- .../xai-grok-shell/src/config/watcher.rs | 944 +++++++++--- .../xai-grok-shell/src/extensions/feedback.rs | 17 +- .../src/extensions/notification.rs | 31 +- .../src/extensions/session_admin.rs | 40 + .../xai-grok-shell/src/sampling/error.rs | 40 +- .../xai-grok-shell/src/session/acp_session.rs | 17 +- .../session/acp_session_impl/auth_retry.rs | 218 +++ .../acp_session_impl/auth_retry_tests.rs | 197 +++ .../src/session/acp_session_impl/recap.rs | 203 ++- .../session/acp_session_impl/sampler_turn.rs | 19 +- .../src/session/acp_session_impl/spawn.rs | 1 + .../session/acp_session_impl/tasks_cancel.rs | 3 + .../src/session/acp_session_impl/turn.rs | 247 ++-- .../src/session/acp_session_impl/types.rs | 35 +- .../auth_error_no_retry_tests.rs | 46 +- .../cancel_running_task_tests.rs | 4 + .../acp_session_tests/idle_resume_tests.rs | 1 + .../inline_auto_compact_flow_tests.rs | 5 + .../acp_session_tests/memory_config_tests.rs | 1 + .../replay_buffer_send_update_tests.rs | 3 + .../src/session/acp_session_tests/support.rs | 1 + .../turn/auth_retry_budget_tests.rs | 308 ++++ .../xai-grok-shell/src/session/commands.rs | 16 +- .../xai-grok-shell/src/session/compaction.rs | 52 +- .../src/session/compaction_config.rs | 114 ++ .../helpers/full_replace_compaction.rs | 5 + .../src/session/helpers/session_compact.rs | 135 +- .../xai-grok-shell/src/session/persistence.rs | 8 + .../src/session/unified_list/mod.rs | 63 +- .../src/terminal/pty_session.rs | 201 ++- .../xai-grok-shell/src/test_support/mod.rs | 9 + .../xai-grok-shell/src/util/dual_clock.rs | 52 + .../src/util/dual_clock_tests.rs | 18 + crates/codegen/xai-grok-shell/src/util/mod.rs | 1 + .../xai-grok-shell/tests/test_leader_soak.rs | 8 +- .../tests/test_registry_churn.rs | 37 +- .../tests/test_sampling_client.rs | 4 +- crates/codegen/xai-grok-telemetry/Cargo.toml | 2 + .../codegen/xai-grok-telemetry/src/events.rs | 49 +- .../xai-grok-telemetry/src/session_ctx.rs | 64 + .../xai-grok-telemetry/src/unified_log.rs | 87 ++ .../xai-grok-test-support/src/mock_server.rs | 185 ++- crates/codegen/xai-grok-tools/Cargo.toml | 4 + crates/codegen/xai-grok-tools/build.rs | 170 +++ .../grok_build/task_output/mod.rs | 57 +- .../task_output/terminal_command.rs | 5 +- .../grok_build/task_output/wait_tasks.rs | 6 +- .../xai-grok-tools/src/registry/types.rs | 151 +- .../src/reminders/task_completion.rs | 4 + .../xai-grok-tools/src/types/context.rs | 158 ++ crates/codegen/xai-grok-version/Cargo.toml | 2 +- .../src/rpc/export.rs | 1 + .../xai-grok-workspace-types/src/rpc/mod.rs | 1 + crates/codegen/xai-grok-workspace/Cargo.toml | 2 + .../xai-grok-workspace/src/activity.rs | 384 ++++- .../codegen/xai-grok-workspace/src/error.rs | 6 + .../codegen/xai-grok-workspace/src/handle.rs | 21 +- .../xai-grok-workspace/src/hub_server.rs | 30 +- crates/codegen/xai-grok-workspace/src/lib.rs | 2 +- .../src/permission/manager.rs | 6 +- .../src/permission/shell_access.rs | 163 ++- .../src/preview_supervisor.rs | 226 ++- .../xai-grok-workspace/src/rpc_envelope.rs | 6 + .../xai-computer-hub-sdk/src/connection.rs | 1296 ++++++++++++++--- .../xai-computer-hub-sdk/src/harness.rs | 4 +- .../xai-computer-hub-sdk/src/metrics.rs | 21 +- .../common/xai-computer-hub-sdk/src/server.rs | 24 +- .../src/code_compaction/failure.rs | 6 + .../src/code_compaction/sample.rs | 20 + .../xai-grok-compaction/src/reminder.rs | 6 + crates/common/xai-tool-protocol/src/frames.rs | 139 ++ crates/common/xai-tool-protocol/src/lib.rs | 30 +- crates/common/xai-tool-types/src/lib.rs | 19 +- crates/common/xai-tool-types/src/task.rs | 74 +- .../src/team_managed_config_types.rs | 83 +- 165 files changed, 15171 insertions(+), 1979 deletions(-) create mode 100644 crates/codegen/xai-grok-pager/src/views/welcome/workspace_mode.rs create mode 100644 crates/codegen/xai-grok-shell/benches/skills_watcher_startup.rs create mode 100644 crates/codegen/xai-grok-shell/changelogs/0.2.117.json create mode 100644 crates/codegen/xai-grok-shell/changelogs/0.2.117.md create mode 100644 crates/codegen/xai-grok-shell/src/agent/mvp_agent/session_registry.rs create mode 100644 crates/codegen/xai-grok-shell/src/session/acp_session_impl/auth_retry.rs create mode 100644 crates/codegen/xai-grok-shell/src/session/acp_session_impl/auth_retry_tests.rs create mode 100644 crates/codegen/xai-grok-shell/src/session/acp_session_tests/turn/auth_retry_budget_tests.rs create mode 100644 crates/codegen/xai-grok-shell/src/util/dual_clock.rs create mode 100644 crates/codegen/xai-grok-shell/src/util/dual_clock_tests.rs create mode 100644 crates/codegen/xai-grok-workspace-types/src/rpc/export.rs diff --git a/Cargo.lock b/Cargo.lock index a8e669b..6e27a66 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,13 +17,24 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher 0.4.4", + "cpufeatures 0.2.17", +] + [[package]] name = "aes" version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "66bd29a732b644c0431c6140f370d097879203d79b80c94a6747ba0872adaef8" dependencies = [ - "cipher", + "cipher 0.5.1", "cpubits", "cpufeatures 0.3.0", ] @@ -1517,6 +1528,25 @@ dependencies = [ "either", ] +[[package]] +name = "bzip2" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49ecfb22d906f800d4fe833b6282cf4dc1c298f5057ca0b5445e5c209735ca47" +dependencies = [ + "bzip2-sys", +] + +[[package]] +name = "bzip2-sys" +version = "0.1.13+1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "225bff33b2141874fe80d71e07d6eec4f85c5c216453dd96388240f96e1acc14" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "cached" version = "0.56.0" @@ -1583,7 +1613,7 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "98db6aeaef0eeef2c1e3ce9a27b739218825dae116076352ac3777076aa22225" dependencies = [ - "cipher", + "cipher 0.5.1", ] [[package]] @@ -1694,6 +1724,16 @@ dependencies = [ "half", ] +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common 0.1.6", + "inout 0.1.4", +] + [[package]] name = "cipher" version = "0.5.1" @@ -1701,7 +1741,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e34d8227fe1ba289043aeb13792056ff80fd6de1a9f49137a5f499de8e8c78ea" dependencies = [ "crypto-common 0.2.1", - "inout", + "inout 0.2.2", ] [[package]] @@ -2338,6 +2378,22 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "ctor" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec09e802f5081de6157da9a75701d6c713d8dc3ba52571fd4bd25f412644e8a6" +dependencies = [ + "ctor-proc-macro", + "dtor", +] + +[[package]] +name = "ctor-proc-macro" +version = "0.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2931af7e13dc045d8e9d26afccc6fa115d64e115c9c84b1166288b46f6782c2" + [[package]] name = "cursor-icon" version = "1.2.0" @@ -2522,6 +2578,12 @@ dependencies = [ "uuid", ] +[[package]] +name = "deflate64" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26bf8fc351c5ed29b5c2f0cbbac1b209b74f60ecd62e675a998df72c49af5204" + [[package]] name = "deltae" version = "0.3.2" @@ -2808,6 +2870,21 @@ dependencies = [ "dtoa", ] +[[package]] +name = "dtor" +version = "0.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97cbdf2ad6846025e8e25df05171abfb30e3ababa12ee0a0e44b9bbe570633a8" +dependencies = [ + "dtor-proc-macro", +] + +[[package]] +name = "dtor-proc-macro" +version = "0.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7454e41ff9012c00d53cf7f475c5e3afa3b91b7c90568495495e8d9bf47a1055" + [[package]] name = "dunce" version = "1.0.5" @@ -5385,6 +5462,15 @@ dependencies = [ "libc", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + [[package]] name = "inout" version = "0.2.2" @@ -5651,6 +5737,15 @@ dependencies = [ "simple_asn1", ] +[[package]] +name = "kamadak-exif" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef4fc70d0ab7e5b6bafa30216a6b48705ea964cdfc29c050f2412295eba58077" +dependencies = [ + "mutate_once", +] + [[package]] name = "kanal" version = "0.1.1" @@ -5967,6 +6062,27 @@ dependencies = [ "url", ] +[[package]] +name = "lzma-rs" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "297e814c836ae64db86b36cf2a557ba54368d03f6afcd7d947c266692f71115e" +dependencies = [ + "byteorder", + "crc", +] + +[[package]] +name = "lzma-sys" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fda04ab3764e6cde78b9974eec4f779acaba7c4e84b36eca3cf77c581b85d27" +dependencies = [ + "cc", + "libc", + "pkg-config", +] + [[package]] name = "mac" version = "0.1.1" @@ -6307,6 +6423,12 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" +[[package]] +name = "mutate_once" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13d2233c9842d08cfe13f9eac96e207ca6a2ea10b80259ebe8ad0268be27d2af" + [[package]] name = "ndk" version = "0.8.0" @@ -7225,13 +7347,23 @@ dependencies = [ "prost-types", ] +[[package]] +name = "pbkdf2" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2" +dependencies = [ + "digest 0.10.7", + "hmac", +] + [[package]] name = "pdf_oxide" version = "0.3.46" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ba152347f68db3aa3fbecc2b0104724d74aad532b186d3d9f4d5541681ddd91" dependencies = [ - "aes", + "aes 0.9.0", "base64", "bitflags 2.13.0", "brotli 8.0.2", @@ -13431,7 +13563,7 @@ dependencies = [ [[package]] name = "xai-grok-pager" -version = "0.2.116" +version = "0.2.117" dependencies = [ "agent-client-protocol", "ansi-to-tui", @@ -13521,7 +13653,7 @@ dependencies = [ [[package]] name = "xai-grok-pager-bin" -version = "0.2.116" +version = "0.2.117" dependencies = [ "anyhow", "clap", @@ -13786,7 +13918,7 @@ dependencies = [ [[package]] name = "xai-grok-shell" -version = "0.2.116" +version = "0.2.117" dependencies = [ "agent-client-protocol", "anyhow", @@ -13802,6 +13934,7 @@ dependencies = [ "chrono", "clap", "criterion", + "ctor", "dashmap", "dhat", "dirs 6.0.0", @@ -14002,6 +14135,7 @@ dependencies = [ "axum", "bytes", "chrono", + "ctor", "dirs 5.0.1", "filetime", "futures-executor", @@ -14089,6 +14223,7 @@ dependencies = [ "dirs 5.0.1", "dunce", "educe", + "encoding_rs", "flate2", "fs2", "futures", @@ -14099,6 +14234,7 @@ dependencies = [ "image", "indexmap", "infer 0.19.0", + "kamadak-exif", "libc", "minijinja", "nix 0.30.1", @@ -14115,6 +14251,7 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_yaml", + "sha2 0.10.9", "shellexpand", "similar", "strip-ansi-escapes", @@ -14127,6 +14264,7 @@ dependencies = [ "tonic", "tracing", "tracing-subscriber", + "unicode-normalization", "url", "uuid", "which", @@ -14193,7 +14331,7 @@ dependencies = [ [[package]] name = "xai-grok-version" -version = "0.2.116" +version = "0.2.117" dependencies = [ "semver", ] @@ -14260,6 +14398,7 @@ dependencies = [ "rustls", "serde", "serde_json", + "sha1", "sha2 0.10.9", "smallvec", "tar", @@ -14308,6 +14447,7 @@ dependencies = [ "xai-tool-types", "xai-tracing", "xai-tty-utils", + "zip", "zstd", ] @@ -14605,6 +14745,15 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec7a2a501ed189703dba8b08142f057e887dfc4b2cc4db2d343ac6376ba3e0b9" +[[package]] +name = "xz2" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "388c44dc09d76f1536602ead6d325eb532f5c122f17782bd57fb47baeeb767e2" +dependencies = [ + "lzma-sys", +] + [[package]] name = "yaml-rust" version = "0.4.5" @@ -14805,11 +14954,25 @@ version = "3.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "12598812502ed0105f607f941c386f43d441e00148fce9dec3ca5ffb0bde9308" dependencies = [ + "aes 0.8.4", "arbitrary", + "bzip2", + "constant_time_eq", "crc32fast", + "deflate64", "flate2", + "getrandom 0.3.4", + "hmac", "indexmap", + "lzma-rs", "memchr", + "pbkdf2", + "sha1", + "time", + "xz2", + "zeroize", + "zopfli", + "zstd", ] [[package]] @@ -14830,6 +14993,18 @@ version = "1.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3ff05f8caa9038894637571ae6b9e29466c1f4f829d26c9b28f869a29cbe3445" +[[package]] +name = "zopfli" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f05cd8797d63865425ff89b5c4a48804f35ba0ce8d125800027ad6017d2b5249" +dependencies = [ + "bumpalo", + "crc32fast", + "log", + "simd-adler32", +] + [[package]] name = "zstd" version = "0.13.3" diff --git a/Cargo.toml b/Cargo.toml index 2932320..cb451e9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -128,6 +128,7 @@ criterion = "0.6" crossbeam = "0.8" crossterm = "0.28" cryptify = "3.2" +ctor = "0.4" dashmap = "6" derive_more = { version = "2", features = ["add", "add_assign", "debug", "deref", "deref_mut", "display", "from", "from_str", "into", "into_iterator", "try_into"] } dhat = "0.3" @@ -135,6 +136,7 @@ dirs = "5.0" documented = "0.9" dunce = "1" educe = "0.6.0" +encoding_rs = "0.8.35" enum_delegate = "0.2" env_logger = "0.11" eventsource-stream = "0.2" @@ -168,6 +170,7 @@ infer = "0.19.0" insta = "1" itertools = "0.14" jsonschema = "0.30.0" +kamadak-exif = "0.5" libc = "0.2" linkify = "0.10" lipsum = "0.9" @@ -224,6 +227,7 @@ serde_json = "1" serde_path_to_error = "0.1" serde_yaml = "0.9" serial_test = "3" +sha1 = "0.10" sha2 = { version = "0.10", features = ["force-soft"] } shlex = "1" signal-hook = "0.3" @@ -263,6 +267,7 @@ tracing-subscriber = { version = "0.3.23", default-features = false, features = ts-rs = "12.0" tui-scrollbar = "0.2" two-face = { version = "0.4", default-features = false, features = ["syntect-fancy"] } +unicode-normalization = "0.1" unicode-segmentation = "1.12.0" unicode-width = "0.2" url = "2" @@ -328,6 +333,7 @@ xai-tool-types = { path = "crates/common/xai-tool-types" } xai-tracing = { path = "crates/common/xai-tracing" } xai-tty-utils = { path = "crates/codegen/xai-tty-utils" } zbus = { version = "5" } +zip = "3" zstd = "0.13" [profile.release] diff --git a/SOURCE_REV b/SOURCE_REV index 6cab3b8..05c8181 100644 --- a/SOURCE_REV +++ b/SOURCE_REV @@ -1 +1 @@ -2a28b4a86cfc4a4c133c35b7fc2a6a9964387c39 +8d69c91f02bcacf01e98d5aebbf2f92547c45738 diff --git a/clippy.toml b/clippy.toml index abab250..a908d0f 100644 --- a/clippy.toml +++ b/clippy.toml @@ -24,4 +24,5 @@ disallowed-methods = [ { path = "tokio::fs::canonicalize", reason = "returns \\\\?\\ verbatim paths on Windows; use xai_grok_tools::util::fs helpers or spawn_blocking + dunce::canonicalize" }, { path = "std::process::Command::spawn", reason = "an unenrolled child outlives its session; use xai_tty_utils::ProcessScope::enroll" }, { path = "tokio::process::Command::spawn", reason = "an unenrolled child outlives its session; use xai_tty_utils::ProcessScope::enroll" }, + { path = "portable_pty::SlavePty::spawn_command", reason = "an unenrolled pty child outlives its session; enroll the shell with xai_tty_utils::ProcessScope::enroll_terminal_pid" }, ] diff --git a/crates/codegen/ptyctl/src/pty.rs b/crates/codegen/ptyctl/src/pty.rs index 63c4ab1..7719069 100644 --- a/crates/codegen/ptyctl/src/pty.rs +++ b/crates/codegen/ptyctl/src/pty.rs @@ -110,6 +110,8 @@ impl PtyHandle { cmd.env("TERM", "xterm-256color"); cmd.env("COLORTERM", "truecolor"); + // Not session-scoped: ptyctl's child is the process it exists to run. + #[allow(clippy::disallowed_methods)] let child = pair .slave .spawn_command(cmd) diff --git a/crates/codegen/xai-grok-http/src/lib.rs b/crates/codegen/xai-grok-http/src/lib.rs index d0c8d68..7925b3c 100644 --- a/crates/codegen/xai-grok-http/src/lib.rs +++ b/crates/codegen/xai-grok-http/src/lib.rs @@ -415,6 +415,40 @@ pub fn error_cause_chain(err: &dyn std::error::Error) -> String { msg } +/// First OS error code in `err`'s `source()` chain (e.g. 104 `ECONNRESET` on +/// Linux, 10054 on Windows), preferring [`std::io::Error::raw_os_error`] and +/// falling back to the `(os error N)` suffix `io::Error`'s `Display` appends. +/// +/// The fallback is load-bearing: a reset during the TLS handshake arrives as a +/// *custom* `io::Error` (kind `Other`, no raw code) whose only record of the +/// code is that suffix, and without it a rustls reset is indistinguishable +/// from an unreachable host. +/// +/// `+ 'static` because `downcast_ref` resolves the type through +/// [`std::any::Any`], whose type ids only exist for `'static` types. +pub fn find_os_error_code(err: &(dyn std::error::Error + 'static)) -> Option { + let mut cur: Option<&(dyn std::error::Error + 'static)> = Some(err); + while let Some(e) = cur { + if let Some(code) = e.downcast_ref::().and_then(|ioe| { + ioe.raw_os_error() + .or_else(|| parse_os_error(&ioe.to_string())) + }) { + return Some(code); + } + cur = e.source(); + } + None +} + +/// Extract `N` from a message ending in `(os error N)`. +fn parse_os_error(msg: &str) -> Option { + msg.rsplit_once("(os error ")? + .1 + .trim_end_matches(')') + .parse() + .ok() +} + /// How a `reqwest` request/send failure should be treated by a retry loop. #[derive(Debug, PartialEq, Eq)] pub enum TransportFailureKind { @@ -601,6 +635,104 @@ mod tests { ); } + #[test] + fn find_os_error_code_walks_source_chain() { + #[derive(Debug)] + struct IoLeaf(std::io::Error); + impl std::fmt::Display for IoLeaf { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "io leaf") + } + } + impl std::error::Error for IoLeaf { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(&self.0) + } + } + + #[derive(Debug)] + struct Wrapper(IoLeaf); + impl std::fmt::Display for Wrapper { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "wrapper") + } + } + impl std::error::Error for Wrapper { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(&self.0) + } + } + + let err = Wrapper(IoLeaf(std::io::Error::from_raw_os_error(104))); + assert_eq!(find_os_error_code(&err), Some(104)); + assert_eq!(find_os_error_code(&std::io::Error::other("no code")), None); + } + + /// A TLS-handshake reset arrives as a custom `io::Error` with no raw code; + /// live reqwest gives the chain `client error (Connect)` → + /// `Connection reset by peer (os error 54)`. + #[test] + fn recovers_code_from_a_custom_io_error() { + let tls_shaped = std::io::Error::other("Connection reset by peer (os error 54)"); + assert_eq!(tls_shaped.raw_os_error(), None, "precondition: no raw code"); + assert_eq!(find_os_error_code(&tls_shaped), Some(54)); + + let windows_shaped = std::io::Error::other( + "An existing connection was forcibly closed by the remote host. (os error 10054)", + ); + assert_eq!(find_os_error_code(&windows_shaped), Some(10054)); + } + + /// Over a real socket: a mid-request reset must classify as `Interrupted` + /// *and* surface the OS code, which is what lets a fleet report tell "peer + /// reset us" from "server unreachable". + /// + /// Lives here rather than in a caller's crate: a `reqwest` client drags + /// rustls into the test binary, which not every caller's tests tolerate. + #[test] + fn real_connection_reset_classifies_as_interrupted_with_os_code() { + // Closing a socket whose receive queue still holds the request emits + // RST instead of FIN. + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let port = listener.local_addr().expect("addr").port(); + std::thread::spawn(move || { + let sock = listener.accept().expect("accept").0; + sock.set_read_timeout(Some(std::time::Duration::from_secs(5))) + .expect("read timeout"); + let _ = sock.peek(&mut [0u8; 64]); + drop(sock); + }); + + let err = reqwest::blocking::Client::new() + .get(format!("http://127.0.0.1:{port}/oauth2/device/code")) + .send() + .expect_err("reset must fail the request"); + + assert_eq!( + TransportFailure::classify(&err).kind, + TransportFailureKind::Interrupted + ); + assert!( + // ECONNRESET: 54 on macOS, 104 on Linux, 10054 on Windows. + matches!(find_os_error_code(&err), Some(54 | 104 | 10054)), + "reset must carry an OS code, got {:?}", + find_os_error_code(&err) + ); + } + + #[test] + fn parse_os_error_ignores_messages_without_a_code() { + assert_eq!( + parse_os_error("connection closed before message completed"), + None + ); + assert_eq!( + parse_os_error("invalid peer certificate (os error oops)"), + None + ); + assert_eq!(parse_os_error("broken pipe (os error 32)"), Some(32)); + } + #[test] fn origin_client_info_from_meta_extracts_identifier_and_version() { let meta = serde_json::json!({ diff --git a/crates/codegen/xai-grok-pager-bin/Cargo.toml b/crates/codegen/xai-grok-pager-bin/Cargo.toml index cedf8c8..cbc0d20 100644 --- a/crates/codegen/xai-grok-pager-bin/Cargo.toml +++ b/crates/codegen/xai-grok-pager-bin/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "xai-grok-pager-bin" -version = "0.2.116" +version = "0.2.117" edition.workspace = true license = "Apache-2.0" authors = ["xAI"] @@ -90,7 +90,9 @@ default = [ default-bazel = [ "jemalloc", "sandbox-enforce", + "local-workspace", ] jemalloc = ["dep:tikv-jemallocator", "dep:tikv-jemalloc-sys", "dep:tikv-jemalloc-ctl"] sandbox-enforce = ["xai-grok-pager/sandbox-enforce"] +local-workspace = ["xai-grok-pager/local-workspace"] release-dist = ["xai-grok-pager/release-dist"] diff --git a/crates/codegen/xai-grok-pager-pty-harness/src/pty.rs b/crates/codegen/xai-grok-pager-pty-harness/src/pty.rs index c8c063f..c917e0f 100644 --- a/crates/codegen/xai-grok-pager-pty-harness/src/pty.rs +++ b/crates/codegen/xai-grok-pager-pty-harness/src/pty.rs @@ -137,6 +137,7 @@ impl PtyController { // portable-pty calls setsid on Unix. Windows Job enrollment is a // best-effort post-spawn attachment, so a very short-lived descendant // may escape before enrollment; diagnostics preserve that downgrade. + #[allow(clippy::disallowed_methods)] let child = pair.slave.spawn_command(cmd)?; #[cfg(unix)] let process_pid = child diff --git a/crates/codegen/xai-grok-pager/Cargo.toml b/crates/codegen/xai-grok-pager/Cargo.toml index 804a7db..6435fce 100644 --- a/crates/codegen/xai-grok-pager/Cargo.toml +++ b/crates/codegen/xai-grok-pager/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "xai-grok-pager" -version = "0.2.116" +version = "0.2.117" edition.workspace = true license = "Apache-2.0" authors = ["xAI"] @@ -290,7 +290,9 @@ default-bazel = [ "jemalloc", "sandbox-enforce", "test-support", + "local-workspace", ] +local-workspace = ["xai-grok-shell/local-workspace"] # No-op on this crate: the actual `#[global_allocator]` (and the # `tikv-jemallocator` dep) live on the composition-root binary # `xai-grok-pager-bin`, which owns `main.rs`. Kept as an empty feature so the diff --git a/crates/codegen/xai-grok-pager/docs/user-guide/03-keyboard-shortcuts.md b/crates/codegen/xai-grok-pager/docs/user-guide/03-keyboard-shortcuts.md index 92afaf2..468ea38 100644 --- a/crates/codegen/xai-grok-pager/docs/user-guide/03-keyboard-shortcuts.md +++ b/crates/codegen/xai-grok-pager/docs/user-guide/03-keyboard-shortcuts.md @@ -253,7 +253,7 @@ Bindings while the [Agent Dashboard](23-dashboard.md) is focused (`Ctrl+\` or `/ | `Ctrl+R` | Rename the selected agent | | `Ctrl+T` | Pin / unpin | | `Ctrl+G` | Toggle grouping (state ↔ working directory) | -| `Ctrl+X` | Stop a running turn, or press twice within 2s to close the session | +| `Ctrl+X` | Cancel a running turn, or press twice within 2s to permanently delete | | `Ctrl+O` | Toggle always-approve on the selected agent | | `Tab` | Toggle focus between the list and the dispatch / peek input | | `Esc` | Step back (cancel search → close peek → clear filter → unfocus → unselect → exit) | diff --git a/crates/codegen/xai-grok-pager/docs/user-guide/04-slash-commands.md b/crates/codegen/xai-grok-pager/docs/user-guide/04-slash-commands.md index 2bc4f81..12d3d5f 100644 --- a/crates/codegen/xai-grok-pager/docs/user-guide/04-slash-commands.md +++ b/crates/codegen/xai-grok-pager/docs/user-guide/04-slash-commands.md @@ -90,7 +90,7 @@ Leave the current session and return to the welcome screen. Alias: `/welcome`. Delete the current session's history and return to the welcome screen. Confirms first. -To delete a session you are not in, open `/resume` and press `d` then `y`. +To delete a session you are not in, open `/resume` or the welcome session list and press `d` then `y`. On the dashboard, press `Ctrl+X` twice or click `[✗]`. ### `/rename` diff --git a/crates/codegen/xai-grok-pager/docs/user-guide/17-sessions.md b/crates/codegen/xai-grok-pager/docs/user-guide/17-sessions.md index cf2f329..f3ddb25 100644 --- a/crates/codegen/xai-grok-pager/docs/user-guide/17-sessions.md +++ b/crates/codegen/xai-grok-pager/docs/user-guide/17-sessions.md @@ -68,7 +68,7 @@ Alias: `/exit`. To leave the current session but stay in Grok, use `/home` to re /delete ``` -Confirms, then permanently removes the session history and returns to the welcome screen. From `/resume`, press `d` then `y` on a row to delete a session you are not currently in. +Confirms, then permanently removes the session history and returns to the welcome screen. From `/resume` or the welcome session list, press `d` then `y`. On the [Agent Dashboard](23-dashboard.md), `Ctrl+X` twice (or hover `[✗]`) permanently deletes. --- diff --git a/crates/codegen/xai-grok-pager/docs/user-guide/23-dashboard.md b/crates/codegen/xai-grok-pager/docs/user-guide/23-dashboard.md index 0245f86..6fb1985 100644 --- a/crates/codegen/xai-grok-pager/docs/user-guide/23-dashboard.md +++ b/crates/codegen/xai-grok-pager/docs/user-guide/23-dashboard.md @@ -89,7 +89,8 @@ The dispatch input uses the same prompt chrome as the agent view. Press | `Ctrl+R` | Rename selected row | | `Ctrl+T` | Pin / unpin | | `Ctrl+G` | Toggle grouping (state ↔ directory) | -| `Ctrl+X` | Stop / kill (two presses within 2s to close a session) | +| `Ctrl+X` | Cancel a running turn, or press twice within 2s to permanently delete | +| Hover + click `[✗]` | Permanently delete an idle/done row (click again to confirm) | | `Shift+↑` / `Shift+↓` | Reorder pinned rows | | `Esc` | Step back: cancel search → close peek → clear filter → unfocus dispatch → unselect row → exit. Never clears a typed dispatch draft (`Ctrl+U` / `Ctrl+C` for that) | | `Ctrl+\` | Return from details view, or exit dashboard | @@ -130,11 +131,14 @@ There is **no** “mark completed” command. Row state is derived from the agen - **Completed** / **Failed** when work ends on its own (turn finished and no background task / monitor / `/loop` still running). - **`Ctrl+X` once** while a turn is running cancels the turn. -- **`Ctrl+X` twice** (within 2s) on an idle / stopped row **closes** the - session and removes it from the live roster. +- **`Ctrl+X` twice** (within 2s) **permanently deletes** the session + (same as `/delete`). Hover an idle/done row to swap age for `[✗]` and + click twice to confirm. - In the details view, `/exit` also closes the session (Esc only returns). + `/delete` inside an attached agent wipes that session and returns home. -Use close/stop when you want a row gone; there is no manual complete flag. +There is no manual complete flag. Use `/exit` to leave a session without +deleting history. --- diff --git a/crates/codegen/xai-grok-pager/src/actions/defaults.rs b/crates/codegen/xai-grok-pager/src/actions/defaults.rs index bd5ba33..d5d3431 100644 --- a/crates/codegen/xai-grok-pager/src/actions/defaults.rs +++ b/crates/codegen/xai-grok-pager/src/actions/defaults.rs @@ -963,8 +963,8 @@ pub(super) fn default_actions( }, ActionDef { id: ActionId::DashboardStop, - label: "stop", - description: "Stop / Close agent", + label: "delete", + description: "Stop / Delete agent", default_key: key!('x', CONTROL), alt_keys: vec![], category: Category::Dashboard, @@ -973,7 +973,7 @@ pub(super) fn default_actions( hint_key_display: None, requires_confirmation: false, long_help: Some( - "Stops the selected agent and removes its row from the dashboard; a running turn is interrupted first.\nUse it to clear finished or unwanted agents without attaching to them.\nThe in-overlay equivalent (Ctrl+X) confirms before stopping.", + "On a busy top-level row, Ctrl+X cancels the running turn. Once the row is idle, press Ctrl+X again within 2s to permanently delete the session.\nOn a subagent row, Ctrl+X kills the subagent.", ), }, ActionDef { diff --git a/crates/codegen/xai-grok-pager/src/app/acp_handler/interactions.rs b/crates/codegen/xai-grok-pager/src/app/acp_handler/interactions.rs index c3eddda..9fc900b 100644 --- a/crates/codegen/xai-grok-pager/src/app/acp_handler/interactions.rs +++ b/crates/codegen/xai-grok-pager/src/app/acp_handler/interactions.rs @@ -56,7 +56,7 @@ pub(crate) fn handle_ask_user_question( // If a question is already active, cancel it before replacing. if let Some(mut old_qv) = agent.question_view.take() { - agent.turn_paused_duration += old_qv.opened_at.elapsed(); + agent.record_question_pause(&old_qv); tracing::warn!( old_tool_call_id = %old_qv.tool_call_id, new_tool_call_id = %ext_req.tool_call_id, diff --git a/crates/codegen/xai-grok-pager/src/app/acp_handler/permissions.rs b/crates/codegen/xai-grok-pager/src/app/acp_handler/permissions.rs index 9993e92..14bea83 100644 --- a/crates/codegen/xai-grok-pager/src/app/acp_handler/permissions.rs +++ b/crates/codegen/xai-grok-pager/src/app/acp_handler/permissions.rs @@ -133,8 +133,17 @@ fn enqueue_permission( let subagent_label = resolve_subagent_label(agent, &perm.request.session_id); // 3. Build title and description from the tool call. - let (title, description, bash_command_raw) = - build_permission_display(&perm.request, bash_highlights.as_ref()); + let (title, description, bash_command_raw) = build_permission_display( + &perm.request, + bash_highlights.as_ref(), + #[cfg(feature = "local-workspace")] + matches!( + agent.workspace_mode, + crate::views::welcome::WelcomeWorkspaceMode::LocalWorkspace + ), + #[cfg(not(feature = "local-workspace"))] + false, + ); // 4. Assign a monotonic ID. let perm_id = agent.next_perm_req_id; @@ -236,6 +245,7 @@ fn resolve_subagent_label(agent: &AgentView, session_id: &acp::SessionId) -> Opt fn build_permission_display( req: &acp::RequestPermissionRequest, bash_highlights: Option<&BashCommandHighlights>, + session_local_workspace: bool, ) -> (String, Vec, Option) { let is_bash = bash_highlights.is_some(); @@ -303,11 +313,29 @@ fn build_permission_display( } }; + let title = qualify_permission_title_for_local_workspace(title, session_local_workspace); let description = permission_description_lines(req); let bash_cmd = if is_execute { raw_command } else { None }; (title, description, bash_cmd) } +/// Per-session HITL copy — not process-global CLI stamp. +fn qualify_permission_title_for_local_workspace( + title: String, + session_local_workspace: bool, +) -> String { + if !session_local_workspace { + return title; + } + if title.contains("on your machine") { + return title; + } + if let Some(stripped) = title.strip_suffix('?') { + return format!("{stripped} (on your machine)?"); + } + format!("{title} (on your machine)") +} + /// Lines shown under the permission title: protected-edit note (if any), then /// MCP planned-argument lines (empty for bash/edit). fn permission_description_lines(req: &acp::RequestPermissionRequest) -> Vec { @@ -450,3 +478,20 @@ pub(super) fn apply_recap_block(agent: &mut AgentView, auto: bool, recap_block: } } } + +#[cfg(all(test, feature = "local-workspace"))] +mod tests { + use super::*; + + #[test] + fn permission_title_qualifies_for_local_workspace() { + assert_eq!( + qualify_permission_title_for_local_workspace("Allow Edit?".into(), false), + "Allow Edit?" + ); + assert_eq!( + qualify_permission_title_for_local_workspace("Allow Edit?".into(), true), + "Allow Edit (on your machine)?" + ); + } +} diff --git a/crates/codegen/xai-grok-pager/src/app/actions.rs b/crates/codegen/xai-grok-pager/src/app/actions.rs index b888d3f..227a41e 100644 --- a/crates/codegen/xai-grok-pager/src/app/actions.rs +++ b/crates/codegen/xai-grok-pager/src/app/actions.rs @@ -129,6 +129,9 @@ pub enum Action { /// load effect; under `--chat`, local Build disk rows are refused in /// dispatch (never coerced). LoadSession(String, Option, bool), + /// Welcome Local workspace ACK confirmed (y); write ack + start session. + #[cfg(feature = "local-workspace")] + ConfirmWelcomeLocalWorkspaceAck, /// Create a new session with a client-chosen session ID (`--session-id`). NewSessionWithId(String), /// Startup `--fork-session`: fork `parent` then load the child. @@ -822,9 +825,12 @@ pub enum Action { DashboardCommitRename, /// Cancel an in-progress rename without committing. DashboardCancelRename, - /// Stop / kill the selected row (top-level: cancel turn → close; - /// subagent: kill). Double-press protected for top-level rows. + /// Ctrl+X on the selected row. Top-level: cancels a running turn on a + /// busy row, else double-press permanently deletes an idle row. + /// Subagent: kills the subagent. DashboardStop, + /// Confirm permanent delete of the armed dashboard row. + DashboardDelete, /// Cycle the dispatch input's mode for the next spawned agent /// (Normal → Plan → Always-Approve → Normal). Bound to Shift+Tab. DashboardCycleMode, @@ -1380,6 +1386,8 @@ pub enum AfterSessionDelete { Stay, /// `/delete` — return to welcome. Welcome, + /// Stay on the dashboard. + Dashboard, } #[derive(Debug)] pub enum Effect { @@ -1470,6 +1478,10 @@ pub enum Effect { /// the response is dropped when no longer current, so out-of-order /// completions can't clobber newer results. seq: u64, + /// Optional unified-list `kind` facet filter (`"chat"` / `"build"`). + /// When set, stamped as `_meta["x.ai/facetFilters"].kind` so the shell + /// honors multi-source history under `--chat` instead of forcing chat-only. + kind_filter: Option>, }, /// Coalesce picker search keystrokes: fires /// [`TaskResult::SessionSearchDebounceExpired`] after a short sleep; the diff --git a/crates/codegen/xai-grok-pager/src/app/agent.rs b/crates/codegen/xai-grok-pager/src/app/agent.rs index a3ddc61..813cc4d 100644 --- a/crates/codegen/xai-grok-pager/src/app/agent.rs +++ b/crates/codegen/xai-grok-pager/src/app/agent.rs @@ -580,6 +580,16 @@ impl AgentState { pub fn is_turn_running(&self) -> bool { matches!(self, Self::TurnRunning) } + /// Manual `/compact` is in flight (stoppable via session/cancel). + pub fn is_compact_running(&self) -> bool { + matches!( + self, + Self::CommandRunning { + command: AgentCommand::Compact, + .. + } + ) + } /// Either a turn or command cancel is in progress. pub fn is_cancelling(&self) -> bool { matches!(self, Self::TurnCancelling | Self::CommandCancelling { .. }) @@ -873,6 +883,18 @@ impl AgentSession { pub fn finish_command(&mut self) { self.state = AgentState::Idle; } + /// Mark an in-flight `/compact` as cancelling (waiting for CompactComplete). + pub fn cancel_compact_command(&mut self) { + if let AgentState::CommandRunning { + command: AgentCommand::Compact, + .. + } = &self.state + { + self.state = AgentState::CommandCancelling { + command: AgentCommand::Compact, + }; + } + } /// Push a prompt onto the back of the queue. Returns the assigned ID. pub fn enqueue_prompt(&mut self, text: String) -> u64 { self.enqueue_entry(text, QueueEntryKind::Prompt) diff --git a/crates/codegen/xai-grok-pager/src/app/agent_view/input.rs b/crates/codegen/xai-grok-pager/src/app/agent_view/input.rs index ca6f0ce..5d90892 100644 --- a/crates/codegen/xai-grok-pager/src/app/agent_view/input.rs +++ b/crates/codegen/xai-grok-pager/src/app/agent_view/input.rs @@ -989,6 +989,7 @@ impl AgentView { } if registry.matches_id(ActionId::CancelTurn, key) && (self.session.state.is_turn_running() + || self.session.state.is_compact_running() || self.session.state.is_cancelling()) { self.dismiss_jump_picker(); @@ -1273,7 +1274,7 @@ impl AgentView { ) -> InputOutcome { match action_id { ActionId::CancelTurn => { - if self.session.state.is_turn_running() { + if self.session.state.is_turn_running() || self.session.state.is_compact_running() { self.cancel_trigger_hint = Some(crate::app::actions::CancelTrigger::CtrlC); return InputOutcome::Action(Action::CancelTurn); } @@ -1541,6 +1542,40 @@ mod background_and_tasks_shortcut_tests { } } #[test] + fn shortcuts_key_tears_down_history_and_opens_cheatsheet() { + use crate::views::modal::ActiveModal; + let registry = ActionRegistry::defaults(); + let history = [HistoryEntry { + text: "earlier prompt".into(), + }]; + for key in [ + KeyEvent::new(KeyCode::Char('.'), KeyModifiers::CONTROL), + KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL), + ] { + for browse in [true, false] { + let mut agent = make_agent(); + agent.set_active_pane(AgentPane::Prompt, true); + if browse { + agent.prompt.history_search.activate_browse(&history, ""); + agent.prompt.set_text("earlier prompt"); + } else { + agent.prompt.history_search.activate(&history, "query"); + agent.prompt.set_text("query"); + } + let out = agent.handle_prompt_key_with_registry_for_test(&key, ®istry); + assert!(matches!(out, InputOutcome::Changed)); + assert!( + matches!(agent.active_modal, Some(ActiveModal::ShortcutsHelp { .. })), + "shortcuts key must open the cheatsheet" + ); + assert!( + !agent.prompt.history_search.is_active(), + "history overlay must be torn down first" + ); + } + } + } + #[test] fn ctrl_b_preempts_file_search_without_mutating_it() { let registry = ActionRegistry::defaults(); let mut agent = make_agent(); @@ -2278,7 +2313,12 @@ mod esc_would_cancel_turn_tests { mod jump_backout_key_tests { use super::test_fixtures::make_agent; use super::{AgentPane, AgentView}; + use crate::actions::ActionRegistry; + use crate::app::actions::Action; + use crate::app::agent::{AgentCommand, AgentState}; + use crate::app::app_view::InputOutcome; use crate::views::jump::{JumpRestore, JumpState}; + use crossterm::event::{Event, KeyCode, KeyEvent, KeyModifiers}; fn open_jump(agent: &mut AgentView) { agent.jump_state = Some(JumpState { entries: Vec::new(), @@ -2290,6 +2330,9 @@ mod jump_backout_key_tests { }, }); } + fn ctrl_c() -> Event { + Event::Key(KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL)) + } /// In the dashboard overlay, a bare Esc backs out via /// `no_esc_consumer_pending`; the open `/jump` picker must count as a /// consumer so Esc dismisses it (restoring the viewport) instead of @@ -2323,6 +2366,26 @@ mod jump_backout_key_tests { "an open /jump picker owns Esc/Left in the overlay back-out" ); } + /// `/jump` must not swallow Ctrl+C while `/compact` is running — same + /// hatch as a running turn. + #[test] + fn jump_picker_ctrl_c_cancels_compact() { + let mut agent = make_agent(); + agent.session.state = AgentState::CommandRunning { + command: AgentCommand::Compact, + started_at: std::time::Instant::now(), + }; + open_jump(&mut agent); + let outcome = agent.handle_input(&ctrl_c(), &ActionRegistry::defaults()); + assert!( + agent.jump_state.is_none(), + "Ctrl+C during /compact must dismiss the jump picker" + ); + assert!( + matches!(outcome, InputOutcome::Action(Action::CancelTurn)), + "Ctrl+C during /compact with /jump open must cancel, got {outcome:?}" + ); + } } #[cfg(test)] mod voice_stop_click_during_plan_review_tests { diff --git a/crates/codegen/xai-grok-pager/src/app/agent_view/interactions.rs b/crates/codegen/xai-grok-pager/src/app/agent_view/interactions.rs index 7f644bf..8ee077c 100644 --- a/crates/codegen/xai-grok-pager/src/app/agent_view/interactions.rs +++ b/crates/codegen/xai-grok-pager/src/app/agent_view/interactions.rs @@ -1056,7 +1056,7 @@ impl AgentView { return self.submit_question_answers(true); } if let Some(qv) = self.question_view.take() { - self.turn_paused_duration += qv.opened_at.elapsed(); + self.record_question_pause(&qv); self.prompt.restore(qv.stashed_prompt); } self.cleanup_question_state(); @@ -1133,7 +1133,7 @@ impl AgentView { let Some(mut qv) = self.question_view.take() else { return InputOutcome::Changed; }; - self.turn_paused_duration += qv.opened_at.elapsed(); + self.record_question_pause(&qv); if let Some(kind) = qv.local_kind.take() { let is_doctor_fix = matches!( kind, @@ -1262,7 +1262,7 @@ impl AgentView { self.question_view = Some(qv); return PeekAnswerOutcome::Advanced; } - self.turn_paused_duration += qv.opened_at.elapsed(); + self.record_question_pause(&qv); let response = qv.build_accepted_response(); qv.send_ext_response(response); self.prompt.restore(qv.stashed_prompt); diff --git a/crates/codegen/xai-grok-pager/src/app/agent_view/mod.rs b/crates/codegen/xai-grok-pager/src/app/agent_view/mod.rs index 0b47c85..6798ec2 100644 --- a/crates/codegen/xai-grok-pager/src/app/agent_view/mod.rs +++ b/crates/codegen/xai-grok-pager/src/app/agent_view/mod.rs @@ -881,6 +881,12 @@ 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, + /// Durable workspace mode for the in-session status indicator (`--chat`). + #[cfg(feature = "local-workspace")] + pub workspace_mode: crate::views::welcome::WelcomeWorkspaceMode, + /// True when CLI/env locked local workspace at startup for this session. + #[cfg(feature = "local-workspace")] + pub workspace_mode_cli_locked: 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. @@ -925,6 +931,11 @@ pub struct AgentView { /// Accumulated duration the turn timer was paused (while the user was /// answering questions via `AskUserQuestion`). Reset when the turn ends. pub turn_paused_duration: std::time::Duration, + /// Wall-clock twin of `turn_paused_duration`: the same pauses measured on + /// the wall clock, which keeps counting through OS suspend while `Instant` + /// does not. Netted against the wall-anchored turn span so a suspend + /// during an open question isn't reported as worked time. + pub turn_paused_wall: std::time::Duration, /// IDs of interjections this client sent and already rendered locally /// (optimistic echo). The shell broadcasts `x.ai/session/interjection` to /// every attached pane; when our own broadcast echoes back carrying an id diff --git a/crates/codegen/xai-grok-pager/src/app/agent_view/prompt.rs b/crates/codegen/xai-grok-pager/src/app/agent_view/prompt.rs index 8ba3f4e..532fc79 100644 --- a/crates/codegen/xai-grok-pager/src/app/agent_view/prompt.rs +++ b/crates/codegen/xai-grok-pager/src/app/agent_view/prompt.rs @@ -116,6 +116,13 @@ impl AgentView { // dropdown state derived from that text would otherwise steal the // arrows mid-browse. if self.prompt.history_search.is_active() { + // Tear down the history overlay before opening the cheatsheet: it + // renders unconditionally and would bleed around the popup, and Esc + // would otherwise silently resume the browse instead of closing help. + if registry.matches_id(ActionId::ShortcutsHelp, key) { + self.close_history_restoring_saved(); + return self.handle_agent_action_with_registry(ActionId::ShortcutsHelp, registry); + } return self.handle_history_search_key(key); } diff --git a/crates/codegen/xai-grok-pager/src/app/agent_view/render.rs b/crates/codegen/xai-grok-pager/src/app/agent_view/render.rs index e5e335c..786c543 100644 --- a/crates/codegen/xai-grok-pager/src/app/agent_view/render.rs +++ b/crates/codegen/xai-grok-pager/src/app/agent_view/render.rs @@ -1326,6 +1326,20 @@ impl AgentView { }) { status.push("mcp", mcp_line); } + #[cfg(feature = "local-workspace")] + if self.chat_kind || self.app_chat_mode { + let label = self + .workspace_mode + .status_label(self.workspace_mode_cli_locked); + let mut mode_style = Style::default().fg(theme.accent_user).bg(theme.bg_base); + if self.workspace_mode_cli_locked { + mode_style = mode_style.add_modifier(ratatui::style::Modifier::DIM); + } + status.push( + "workspace_mode", + Line::from(Span::styled(label, mode_style)), + ); + } let ctx_used = self.context_state.as_ref().map(|c| c.used); let model_window = self.session.models.get_context_window(); let ctx_total = self diff --git a/crates/codegen/xai-grok-pager/src/app/agent_view/session.rs b/crates/codegen/xai-grok-pager/src/app/agent_view/session.rs index 9a7e1f1..8cd245d 100644 --- a/crates/codegen/xai-grok-pager/src/app/agent_view/session.rs +++ b/crates/codegen/xai-grok-pager/src/app/agent_view/session.rs @@ -117,6 +117,10 @@ impl AgentView { context_state: None, chat_kind: false, app_chat_mode: false, + #[cfg(feature = "local-workspace")] + workspace_mode: crate::views::welcome::WelcomeWorkspaceMode::Sandbox, + #[cfg(feature = "local-workspace")] + workspace_mode_cli_locked: false, credit_balance: None, auto_topup: None, goal_state: None, @@ -134,6 +138,7 @@ impl AgentView { turn_started_at: None, first_activity_logged_for: None, turn_paused_duration: std::time::Duration::ZERO, + turn_paused_wall: std::time::Duration::ZERO, self_interjection_ids: std::collections::HashSet::new(), last_active_at: Some(Instant::now()), current_branch: None, @@ -350,17 +355,35 @@ impl AgentView { child_view.mark_as_subagent_view(); self.subagent_views.insert(child_sid, child_view); } - /// Clear `turn_started_at` and stamp `last_active_at` to "now". + /// Clear the turn-timing fields and stamp `last_active_at` to "now". /// /// Call this from every site that ends a turn (success, failure, - /// cancellation, reconnect cleanup). Centralised so the two - /// fields cannot drift apart at the ~10 termination call sites - /// across `dispatch.rs` and `event_loop.rs`. + /// cancellation, reconnect cleanup). Centralised so the fields cannot + /// drift apart at the ~10 termination call sites across `dispatch.rs` + /// and `event_loop.rs`. The wall anchor is cleared so a later turn that + /// reuses a prompt id (stash-and-resubmit after `/login`) can never + /// wall-max against a previous attempt's anchor in + /// [`honest_turn_elapsed`]. pub fn mark_turn_finished(&mut self) { self.turn_started_at = None; self.turn_paused_duration = std::time::Duration::ZERO; + self.turn_paused_wall = std::time::Duration::ZERO; + self.turn_start_ms = None; + self.turn_start_ms_prompt = None; self.last_active_at = Some(Instant::now()); } + /// Absorb a closing/replaced question view's open span into the turn's + /// pause totals, on both clocks — a close site that updated only the + /// `Instant` pause would resurface suspend time as worked time in + /// [`honest_turn_elapsed`]. + pub(crate) fn record_question_pause( + &mut self, + qv: &crate::views::question_view::QuestionViewState, + ) { + self.turn_paused_duration += qv.opened_at.elapsed(); + self.turn_paused_wall += + wall_since_ms(qv.opened_at_wall_ms, chrono::Utc::now().timestamp_millis()); + } /// Invalidate and clear a minimal `/btw` lifecycle at a session boundary. pub(crate) fn clear_minimal_btw_lifecycle(&mut self) { crate::minimal_api::clear_minimal_btw(self); @@ -637,18 +660,26 @@ impl AgentView { self.reset_follow_ups_for_reload(); dropped_heavy } - /// Effective turn elapsed time, excluding time spent in question views. - /// - /// Subtracts both the accumulated `turn_paused_duration` (from previously - /// closed question views) and the time elapsed since the current question - /// view opened (if one is active). + /// Effective turn elapsed time, excluding time spent in question views + /// (accumulated pauses plus the currently open one, on both clocks). pub fn turn_elapsed(&self) -> Option { - let raw = self.turn_started_at?.elapsed(); - let mut paused = self.turn_paused_duration; + let instant_elapsed = self.turn_started_at?.elapsed(); + let now_ms = chrono::Utc::now().timestamp_millis(); + let mut instant_paused = self.turn_paused_duration; + let mut wall_paused = self.turn_paused_wall; if let Some(qv) = &self.question_view { - paused += qv.opened_at.elapsed(); + instant_paused += qv.opened_at.elapsed(); + wall_paused += wall_since_ms(qv.opened_at_wall_ms, now_ms); } - Some(raw.saturating_sub(paused)) + Some(honest_turn_elapsed(TurnElapsedParams { + instant_elapsed, + instant_paused, + wall_anchor_ms: self.turn_start_ms, + wall_paused, + anchor_prompt: self.turn_start_ms_prompt.as_deref(), + current_prompt: self.session.current_prompt_id.as_deref(), + now_ms, + })) } /// Turn activity for the status spinner, with the implicit "no activity" /// gap during a running inference turn resolved into an explicit @@ -979,6 +1010,168 @@ impl AgentView { self.prompt.set_voice_visible(available); } } +/// Inputs for [`honest_turn_elapsed`]: the turn span and pause total measured +/// on each clock, plus the wire anchor's provenance. `now_ms` is injected so +/// tests control the wall clock. +struct TurnElapsedParams<'a> { + instant_elapsed: std::time::Duration, + instant_paused: std::time::Duration, + /// `turnStartMs` wire anchor (UTC ms) and the prompt id it was stamped + /// for; the anchor counts only when that id matches the running prompt + /// (interleaved deltas can re-stamp it with another prompt's anchor). + wall_anchor_ms: Option, + wall_paused: std::time::Duration, + anchor_prompt: Option<&'a str>, + current_prompt: Option<&'a str>, + now_ms: i64, +} +/// Turn elapsed for [`AgentView::turn_elapsed`], honest across OS suspends +/// (`Instant` pauses while the machine sleeps; the wall clock keeps +/// counting). Each span is netted against pauses measured on its own clock, +/// and the larger net wins; the tests below enumerate the guard cases. +fn honest_turn_elapsed(params: TurnElapsedParams<'_>) -> std::time::Duration { + let instant_net = params.instant_elapsed.saturating_sub(params.instant_paused); + let (Some(start_ms), Some(anchor_prompt), Some(current_prompt)) = ( + params.wall_anchor_ms, + params.anchor_prompt, + params.current_prompt, + ) else { + return instant_net; + }; + if anchor_prompt != current_prompt { + return instant_net; + } + let wall_net = wall_since_ms(start_ms, params.now_ms).saturating_sub(params.wall_paused); + instant_net.max(wall_net) +} +/// Wall-clock span since `start_ms`, clamped to zero when `start_ms` +/// postdates `now_ms` (skew) so a wall span can never go negative. +fn wall_since_ms(start_ms: i64, now_ms: i64) -> std::time::Duration { + std::time::Duration::from_millis(u64::try_from(now_ms.saturating_sub(start_ms)).unwrap_or(0)) +} +#[cfg(test)] +mod honest_turn_elapsed_tests { + use super::*; + use std::time::Duration; + const NOW_MS: i64 = 1_700_000_000_000; + const MIN: u64 = 60; + const HOUR: u64 = 3_600; + /// Valid same-prompt anchor context with zero spans; tests override the + /// fields under test via struct-update syntax. + fn base() -> TurnElapsedParams<'static> { + TurnElapsedParams { + instant_elapsed: Duration::ZERO, + instant_paused: Duration::ZERO, + wall_anchor_ms: None, + wall_paused: Duration::ZERO, + anchor_prompt: Some("p1"), + current_prompt: Some("p1"), + now_ms: NOW_MS, + } + } + #[test] + fn no_wall_anchor_keeps_instant_net() { + assert_eq!( + honest_turn_elapsed(TurnElapsedParams { + instant_elapsed: Duration::from_secs(5 * MIN), + instant_paused: Duration::from_secs(MIN), + ..base() + }), + Duration::from_secs(4 * MIN) + ); + } + #[test] + fn suspend_outside_questions_defers_to_wall_net() { + assert_eq!( + honest_turn_elapsed(TurnElapsedParams { + instant_elapsed: Duration::from_secs(4 * MIN), + wall_anchor_ms: Some(NOW_MS - 2 * HOUR as i64 * 1_000), + ..base() + }), + Duration::from_secs(2 * HOUR) + ); + } + #[test] + fn suspend_while_question_open_is_not_worked_time() { + assert_eq!( + honest_turn_elapsed(TurnElapsedParams { + instant_elapsed: Duration::from_secs(10 * MIN), + instant_paused: Duration::from_secs(5 * MIN), + wall_anchor_ms: Some(NOW_MS - (2 * HOUR as i64 + 10 * MIN as i64) * 1_000), + wall_paused: Duration::from_secs(2 * HOUR + 5 * MIN), + ..base() + }), + Duration::from_secs(5 * MIN) + ); + } + #[test] + fn instant_net_bounds_below_after_backward_wall_jump() { + assert_eq!( + honest_turn_elapsed(TurnElapsedParams { + instant_elapsed: Duration::from_secs(5 * MIN), + wall_anchor_ms: Some(NOW_MS - 1_000), + ..base() + }), + Duration::from_secs(5 * MIN) + ); + } + #[test] + fn foreign_prompt_anchor_falls_back_to_instant_net() { + assert_eq!( + honest_turn_elapsed(TurnElapsedParams { + instant_elapsed: Duration::from_secs(10 * MIN), + instant_paused: Duration::from_secs(4 * MIN), + wall_anchor_ms: Some(NOW_MS - 2 * HOUR as i64 * 1_000), + anchor_prompt: Some("p-other"), + ..base() + }), + Duration::from_secs(6 * MIN) + ); + } + #[test] + fn missing_current_prompt_ignores_anchor() { + assert_eq!( + honest_turn_elapsed(TurnElapsedParams { + instant_elapsed: Duration::from_secs(MIN), + wall_anchor_ms: Some(NOW_MS - 2 * HOUR as i64 * 1_000), + current_prompt: None, + ..base() + }), + Duration::from_secs(MIN) + ); + } + #[test] + fn future_wall_anchor_is_ignored() { + assert_eq!( + honest_turn_elapsed(TurnElapsedParams { + instant_elapsed: Duration::from_secs(MIN), + wall_anchor_ms: Some(NOW_MS + 60_000), + ..base() + }), + Duration::from_secs(MIN) + ); + } + #[test] + fn turn_elapsed_reflects_wall_span_for_current_prompt() { + let mut view = test_agent_view(Some("s1"), std::path::PathBuf::from("/tmp")); + view.turn_started_at = Some(Instant::now()); + view.turn_start_ms = Some(chrono::Utc::now().timestamp_millis() - 60_000); + view.turn_start_ms_prompt = Some("p1".to_string()); + view.session.current_prompt_id = Some("p1".to_string()); + assert!(view.turn_elapsed().unwrap() >= Duration::from_secs(59)); + } + #[test] + fn turn_elapsed_nets_wall_pauses_against_wall_span() { + let mut view = test_agent_view(Some("s1"), std::path::PathBuf::from("/tmp")); + view.turn_started_at = Some(Instant::now()); + view.turn_start_ms = Some(chrono::Utc::now().timestamp_millis() - 60_000); + view.turn_start_ms_prompt = Some("p1".to_string()); + view.session.current_prompt_id = Some("p1".to_string()); + view.turn_paused_wall = Duration::from_secs(45); + let elapsed = view.turn_elapsed().unwrap(); + assert!(elapsed >= Duration::from_secs(14) && elapsed <= Duration::from_secs(16)); + } +} #[cfg(test)] mod resolve_turn_activity_tests { use super::*; diff --git a/crates/codegen/xai-grok-pager/src/app/app_view.rs b/crates/codegen/xai-grok-pager/src/app/app_view.rs index 87a15a9..f696eeb 100644 --- a/crates/codegen/xai-grok-pager/src/app/app_view.rs +++ b/crates/codegen/xai-grok-pager/src/app/app_view.rs @@ -499,7 +499,7 @@ impl PendingAction { } /// Like [`Self::new`] but with an explicit confirm window. Used by /// the dashboard-overlay stop (Ctrl+X), which mirrors the - /// dashboard's [`crate::views::dashboard::state::STOP_CONFIRM_WINDOW`] + /// dashboard's [`crate::views::dashboard::state::CONFIRM_WINDOW`] /// rather than the default double-press TTL. pub fn with_ttl( action: Action, @@ -865,6 +865,12 @@ pub struct AppView { pub welcome_privacy_banner_opt_out_rect: Option, pub welcome_privacy_banner_terms_rect: Option, pub welcome_privacy_banner_policy_rect: Option, + /// Hit-test rects for the welcome workspace-mode picker. + #[cfg(feature = "local-workspace")] + pub welcome_workspace_mode_rects: crate::views::welcome::WorkspaceModeHitRects, + /// Sticky hover flag for the workspace-mode picker (redraw on enter/leave). + #[cfg(feature = "local-workspace")] + pub welcome_on_workspace_mode: bool, /// Transient welcome toast: (message, wall-clock expiry). pub welcome_toast: Option<(String, std::time::Instant)>, /// Sticky hover flag for the privacy banner buttons (redraw on enter/leave). @@ -918,6 +924,7 @@ pub struct AppView { /// [`crate::views::session_picker::effective_filter_query`], skips the /// local fuzzy re-filter for server search results. pub session_picker_entries_query: Option, + pub session_picker_pending_delete: Option, /// Tick counter for welcome screen spinner animation. pub welcome_tick: u64, /// Last shimmer frame drawn on the welcome screen. Lets `tick` throttle the @@ -964,6 +971,22 @@ pub struct AppView { /// profiles on create/load while set. `/chat` does **not** set this /// (uses [`Self::deferred_startup`] one-shot state instead). pub chat_mode: bool, + /// Welcome picker mode; ignored when `local_workspace_startup_locked`. + #[cfg(feature = "local-workspace")] + pub welcome_workspace_mode: crate::views::welcome::WelcomeWorkspaceMode, + /// CLI/env already stamped local workspace; welcome must not override. + #[cfg(feature = "local-workspace")] + pub local_workspace_startup_locked: bool, + /// One-shot next-session stamp: `Some(None)` sandbox, `Some(cfg)` local. + #[cfg(feature = "local-workspace")] + pub welcome_session_local_workspace: + Option>, + /// First-run Local ACK still pending in the TUI. + #[cfg(feature = "local-workspace")] + pub welcome_local_workspace_ack_pending: bool, + /// Next welcome history load is local-disk/build (does not set `chat_mode`). + #[cfg(feature = "local-workspace")] + pub welcome_history_load_as_build: bool, /// Whether mouse capture is currently enabled. Disabled during the /// Authenticating state so the terminal handles native text selection. pub mouse_captured: bool, @@ -1442,6 +1465,10 @@ impl AppView { welcome_privacy_banner_opt_out_rect: None, welcome_privacy_banner_terms_rect: None, welcome_privacy_banner_policy_rect: None, + #[cfg(feature = "local-workspace")] + welcome_workspace_mode_rects: Default::default(), + #[cfg(feature = "local-workspace")] + welcome_on_workspace_mode: false, welcome_toast: None, welcome_on_privacy_banner: false, welcome_on_upgrade_cta: false, @@ -1465,6 +1492,7 @@ impl AppView { session_picker_lanes: Default::default(), session_picker_detail_generation: 0, session_picker_entries_query: None, + session_picker_pending_delete: None, welcome_tick: 0, welcome_shimmer_frame: 0, cli_model_override: None, @@ -1480,6 +1508,16 @@ impl AppView { subagents: false, ask_user: false, chat_mode: false, + #[cfg(feature = "local-workspace")] + welcome_workspace_mode: crate::views::welcome::WelcomeWorkspaceMode::Sandbox, + #[cfg(feature = "local-workspace")] + local_workspace_startup_locked: false, + #[cfg(feature = "local-workspace")] + welcome_session_local_workspace: None, + #[cfg(feature = "local-workspace")] + welcome_local_workspace_ack_pending: false, + #[cfg(feature = "local-workspace")] + welcome_history_load_as_build: false, mouse_captured: true, new_worktree_dialog: None, contextual_hints: Default::default(), @@ -2429,6 +2467,8 @@ impl AppView { self.session_picker_loading, &self.session_picker_lanes, ); + #[cfg(feature = "local-workspace")] + let session_picker_open = self.session_picker_entries.is_some() || sp_loading; let outcome = match self.active_view { ActiveView::Welcome => handle_welcome_input( ev, @@ -2492,7 +2532,24 @@ impl AppView { cwd_has_git_ancestor: self.cwd_has_git_ancestor, session_picker_grouped: self.session_picker_grouped, sp_source_filter: &mut self.session_picker_source_filter, + sp_pending_delete: &mut self.session_picker_pending_delete, chat_mode: self.chat_mode, + #[cfg(feature = "local-workspace")] + workspace_mode: &mut self.welcome_workspace_mode, + #[cfg(feature = "local-workspace")] + workspace_mode_rects: &self.welcome_workspace_mode_rects, + #[cfg(feature = "local-workspace")] + on_workspace_mode: &mut self.welcome_on_workspace_mode, + #[cfg(feature = "local-workspace")] + workspace_mode_startup_locked: self.local_workspace_startup_locked, + #[cfg(feature = "local-workspace")] + workspace_mode_ack_pending: &mut self.welcome_local_workspace_ack_pending, + #[cfg(feature = "local-workspace")] + history_load_as_build: &mut self.welcome_history_load_as_build, + #[cfg(feature = "local-workspace")] + deferred_startup: &mut self.deferred_startup, + #[cfg(feature = "local-workspace")] + session_picker_open, }, ), ActiveView::Agent(id) => { @@ -2537,18 +2594,17 @@ impl AppView { return InputOutcome::Action(Action::DashboardOverlayNext); } Some(crate::actions::ActionId::DashboardOverlayStop) => { - if self - .agents - .get(&id) - .is_some_and(|a| a.session.state.is_turn_running()) - { + if self.agents.get(&id).is_some_and(|a| { + a.session.state.is_turn_running() + || a.session.state.is_compact_running() + }) { return InputOutcome::Action(Action::CancelTurn); } self.pending_action = Some(PendingAction::with_ttl( Action::DashboardOverlayStop, KeyShortcut::from(*key), Some("close this session"), - crate::views::dashboard::state::STOP_CONFIRM_WINDOW, + crate::views::dashboard::state::CONFIRM_WINDOW, )); return InputOutcome::Changed; } @@ -3109,9 +3165,26 @@ struct WelcomeInputCtx<'a> { cwd_has_git_ancestor: bool, session_picker_grouped: bool, sp_source_filter: &'a mut crate::views::session_picker::SourceFilter, + sp_pending_delete: &'a mut Option, /// Process-wide `--chat`: the session picker hides its source filter /// (conversations-only list), so `f` must not cycle it. chat_mode: bool, + #[cfg(feature = "local-workspace")] + workspace_mode: &'a mut crate::views::welcome::WelcomeWorkspaceMode, + #[cfg(feature = "local-workspace")] + workspace_mode_rects: &'a crate::views::welcome::WorkspaceModeHitRects, + #[cfg(feature = "local-workspace")] + on_workspace_mode: &'a mut bool, + #[cfg(feature = "local-workspace")] + workspace_mode_startup_locked: bool, + #[cfg(feature = "local-workspace")] + workspace_mode_ack_pending: &'a mut bool, + #[cfg(feature = "local-workspace")] + history_load_as_build: &'a mut bool, + #[cfg(feature = "local-workspace")] + deferred_startup: &'a mut crate::app::session_startup::DeferredStartupActions, + #[cfg(feature = "local-workspace")] + session_picker_open: bool, } /// Welcome view input -- auth-state-aware routing. fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutcome { @@ -3252,6 +3325,83 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco } return InputOutcome::Unchanged; } + #[cfg(feature = "local-workspace")] + if *ctx.workspace_mode_ack_pending + && matches!(ctx.auth_state, AuthState::Done) + && ctx.has_access + && !ctx.is_zdr_blocked + { + if let Event::Key(key) = ev { + if key.kind == KeyEventKind::Release { + return InputOutcome::Unchanged; + } + if key!('y').matches(key) || key!('Y').matches(key) || key!(Enter).matches(key) { + return InputOutcome::Action(Action::ConfirmWelcomeLocalWorkspaceAck); + } + if key!('n').matches(key) || key!('N').matches(key) || key!(Esc).matches(key) { + *ctx.workspace_mode_ack_pending = false; + *ctx.workspace_mode = crate::views::welcome::WelcomeWorkspaceMode::Sandbox; + let was_worktree = ctx.deferred_startup.worktree; + ctx.deferred_startup.worktree = false; + ctx.deferred_startup.worktree_label = None; + ctx.deferred_startup.worktree_ref = None; + if was_worktree { + ctx.deferred_startup.session = None; + ctx.deferred_startup.preferred_session_id = None; + } + *ctx.history_load_as_build = false; + ctx.deferred_startup.history_load_as_build = false; + crate::views::welcome::workspace_mode::log_welcome_ack("cancelled"); + return InputOutcome::Changed; + } + return InputOutcome::Unchanged; + } + if matches!(ev, Event::Resize(_, _)) { + return InputOutcome::Changed; + } + return InputOutcome::Unchanged; + } + #[cfg(feature = "local-workspace")] + if crate::views::welcome::workspace_mode::picker_interactive( + ctx.chat_mode, + ctx.has_access, + matches!(ctx.auth_state, AuthState::Done), + ctx.is_zdr_blocked, + ctx.session_picker_open, + ctx.workspace_mode_startup_locked, + ) { + if let Event::Key(key) = ev + && key.kind != KeyEventKind::Release + && key!('e', CONTROL).matches(key) + { + *ctx.workspace_mode = ctx.workspace_mode.cycle_next(); + crate::views::welcome::workspace_mode::log_welcome_mode_selected( + *ctx.workspace_mode, + "ctrl_e", + ctx.workspace_mode_startup_locked, + ); + return InputOutcome::Changed; + } + if let Event::Mouse(mouse) = ev + && matches!( + mouse.kind, + crossterm::event::MouseEventKind::Down(crossterm::event::MouseButton::Left) + ) + && let Some(mode) = crate::views::welcome::hit_test_workspace_mode( + ctx.workspace_mode_rects, + mouse.column, + mouse.row, + ) + { + *ctx.workspace_mode = mode; + crate::views::welcome::workspace_mode::log_welcome_mode_selected( + mode, + "click", + ctx.workspace_mode_startup_locked, + ); + return InputOutcome::Changed; + } + } if (ctx.sp_entries.is_some() || ctx.sp_loading) && matches!(ctx.auth_state, AuthState::Done) { use crate::views::picker::{PickerConfig, PickerOutcome, handle_picker_input}; let source_filter = *ctx.sp_source_filter; @@ -3271,6 +3421,19 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco ); let entry_count = entry_map.len(); let non_selectable_flags: Vec = entry_map.iter().map(|e| e.is_none()).collect(); + let focused_is_foreign = match entry_map + .get(ctx.sp_state.selected) + .and_then(|entry| entry.as_ref()) + { + Some(PickerItem::Fuzzy { original_index }) => ctx + .sp_entries + .as_ref() + .and_then(|entries| entries.get(*original_index)) + .is_some_and(|entry| { + crate::app::foreign_sessions::is_foreign_picker_source(&entry.source) + }), + _ => false, + }; let config = PickerConfig { title: Some("Resume session"), show_search_hint: true, @@ -3287,12 +3450,30 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco filter_key_hint: (!ctx.chat_mode).then_some("f"), filter_active: !ctx.chat_mode && source_filter.is_active(), header_note: None, - action_keys: &[], + action_keys: if ctx.chat_mode || focused_is_foreign { + &[] + } else { + &[('d', "delete")] + }, disable_search: false, compact_bottom_bar: false, search_only_on_slash: false, vim_normal_first: crate::appearance::cache::load_vim_mode(), }; + match crate::views::session_picker::handle_pending_delete_key(ctx.sp_pending_delete, ev) { + crate::views::session_picker::PendingDeleteKey::Confirm(pd) => { + return InputOutcome::Action(Action::DeleteSession { + source: pd.source, + session_id: pd.session_id, + cwd: pd.cwd, + }); + } + crate::views::session_picker::PendingDeleteKey::Cancel => { + return InputOutcome::Changed; + } + crate::views::session_picker::PendingDeleteKey::Disarmed + | crate::views::session_picker::PendingDeleteKey::NotArmed => {} + } if let Event::Key(key) = ev { if key.kind == KeyEventKind::Press && (key!('c', CONTROL).matches(key) || key!('d', CONTROL).matches(key)) @@ -3320,7 +3501,11 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco }); } } + let selected_before = ctx.sp_state.selected; let outcome = handle_picker_input(ev, ctx.sp_state, entry_count, &config); + if ctx.sp_pending_delete.is_some() && ctx.sp_state.selected != selected_before { + *ctx.sp_pending_delete = None; + } match outcome { PickerOutcome::Selected(i) => match entry_map.get(i).and_then(|e| e.as_ref()) { Some(PickerItem::Fuzzy { original_index }) => { @@ -3350,6 +3535,7 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco *ctx.sp_entries = None; ctx.sp_state.reset(); *ctx.sp_source_filter = crate::views::session_picker::SourceFilter::default(); + *ctx.sp_pending_delete = None; return InputOutcome::Action(Action::SessionPickerClosed); } PickerOutcome::Expand(i) => { @@ -3443,6 +3629,16 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco PickerOutcome::FilterCycled => { return InputOutcome::Action(Action::CycleSessionSourceFilter); } + PickerOutcome::Action('d') => { + *ctx.sp_pending_delete = + crate::views::session_picker::pending_delete_from_selection( + ctx.sp_state.selected, + &entry_map, + ctx.sp_entries.as_deref(), + ctx.sp_content_results.as_deref(), + ); + return InputOutcome::Changed; + } PickerOutcome::NonSelectableClick(_) | PickerOutcome::TabChanged(_) | PickerOutcome::Action(_) => { @@ -3801,6 +3997,20 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco *ctx.on_upgrade_cta = over_upgrade; return InputOutcome::Changed; } + #[cfg(feature = "local-workspace")] + { + let over_ws = ctx + .workspace_mode_rects + .row + .is_some_and(|r| r.contains(pos)); + if over_ws != *ctx.on_workspace_mode { + *ctx.on_workspace_mode = over_ws; + return InputOutcome::Changed; + } + if over_ws { + return InputOutcome::Changed; + } + } let over_banner = ctx .privacy_banner_opt_in_rect .is_some_and(|r| r.contains(pos)) @@ -4332,6 +4542,9 @@ impl AppView { subscription_tier: self.subscription_tier.as_deref(), session_picker_grouped: self.session_picker_grouped, session_picker_source_filter: self.session_picker_source_filter, + session_picker_pending_delete: self + .session_picker_pending_delete + .is_some(), chat_mode: self.chat_mode, credit_balance: self.credit_balance.as_ref(), auto_topup: self.auto_topup.as_ref(), @@ -4342,6 +4555,12 @@ impl AppView { welcome_announcement_expanded: self.welcome_announcement.expanded, upgrade_cta: hero_cta.map(|(_owner, label, _)| label), privacy_banner, + #[cfg(feature = "local-workspace")] + workspace_mode: self.welcome_workspace_mode, + #[cfg(feature = "local-workspace")] + workspace_mode_startup_locked: self.local_workspace_startup_locked, + #[cfg(feature = "local-workspace")] + workspace_mode_ack_pending: self.welcome_local_workspace_ack_pending, }; let result = crate::views::welcome::render_welcome( view_area, @@ -4364,6 +4583,10 @@ impl AppView { result.privacy_banner_opt_out_rect; self.welcome_privacy_banner_terms_rect = result.privacy_banner_terms_rect; self.welcome_privacy_banner_policy_rect = result.privacy_banner_policy_rect; + #[cfg(feature = "local-workspace")] + { + self.welcome_workspace_mode_rects = result.workspace_mode_rects; + } self.welcome_changelog_cta_rect = result.changelog_cta_rect; if let Some((ref msg, _)) = self.welcome_toast { crate::views::welcome::paint_welcome_toast( @@ -5728,6 +5951,16 @@ pub(crate) mod tests { subagents: false, ask_user: false, chat_mode: false, + #[cfg(feature = "local-workspace")] + welcome_workspace_mode: crate::views::welcome::WelcomeWorkspaceMode::Sandbox, + #[cfg(feature = "local-workspace")] + local_workspace_startup_locked: false, + #[cfg(feature = "local-workspace")] + welcome_session_local_workspace: None, + #[cfg(feature = "local-workspace")] + welcome_local_workspace_ack_pending: false, + #[cfg(feature = "local-workspace")] + welcome_history_load_as_build: false, mouse_captured: true, new_worktree_dialog: None, contextual_hints: Default::default(), @@ -5812,6 +6045,10 @@ pub(crate) mod tests { welcome_privacy_banner_opt_out_rect: None, welcome_privacy_banner_terms_rect: None, welcome_privacy_banner_policy_rect: None, + #[cfg(feature = "local-workspace")] + welcome_workspace_mode_rects: Default::default(), + #[cfg(feature = "local-workspace")] + welcome_on_workspace_mode: false, welcome_toast: None, welcome_on_privacy_banner: false, welcome_on_upgrade_cta: false, @@ -5835,6 +6072,7 @@ pub(crate) mod tests { session_picker_lanes: Default::default(), session_picker_detail_generation: 0, session_picker_entries_query: None, + session_picker_pending_delete: None, welcome_tick: 0, welcome_shimmer_frame: 0, startup_warnings: Vec::new(), @@ -11405,6 +11643,25 @@ pub(crate) mod tests { "Ctrl+X must be intercepted before the agent sees it", ); } + /// Overlay Ctrl+X during `/compact` cancels compaction (same as `[stop]`). + #[test] + fn overlay_ctrl_x_compact_running_cancels_without_arming() { + use crate::app::agent::{AgentCommand, AgentState}; + let (mut app, id) = neutral_overlay_app(); + app.agents.get_mut(&id).unwrap().session.state = AgentState::CommandRunning { + command: AgentCommand::Compact, + started_at: std::time::Instant::now(), + }; + let outcome = app.handle_input(&key_event(KeyCode::Char('x'), KeyModifiers::CONTROL)); + assert!( + matches!(outcome, InputOutcome::Action(Action::CancelTurn)), + "Ctrl+X during /compact must cancel, got {outcome:?}", + ); + assert!( + app.pending_action.is_none(), + "Ctrl+X during /compact must not arm close confirm", + ); + } /// Overlay Ctrl+X on a non-turn busy agent (command in flight, /// cancel pending) — `Action::CancelTurn` would no-op for these /// states, so the press arms the two-press close instead of @@ -11413,10 +11670,6 @@ pub(crate) mod tests { fn overlay_ctrl_x_command_or_cancelling_agent_arms_close_confirm() { use crate::app::agent::{AgentCommand, AgentState}; let states = [ - AgentState::CommandRunning { - command: AgentCommand::Compact, - started_at: std::time::Instant::now(), - }, AgentState::TurnCancelling, AgentState::CommandCancelling { command: AgentCommand::Compact, @@ -11739,4 +11992,183 @@ pub(crate) mod tests { InputOutcome::Action(Action::CycleSessionSourceFilter) )); } + #[cfg(feature = "local-workspace")] + #[test] + fn welcome_ctrl_e_cycles_workspace_mode() { + use crate::views::welcome::WelcomeWorkspaceMode; + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.auth_state = AuthState::Done; + app.trust_state = TrustState::Done; + assert_eq!(app.welcome_workspace_mode, WelcomeWorkspaceMode::Sandbox); + let key = Event::Key(KeyEvent::new(KeyCode::Char('e'), KeyModifiers::CONTROL)); + let outcome = app.handle_input(&key); + assert!(matches!(outcome, InputOutcome::Changed)); + assert_eq!( + app.welcome_workspace_mode, + WelcomeWorkspaceMode::LocalWorkspace + ); + let _ = app.handle_input(&key); + assert_eq!(app.welcome_workspace_mode, WelcomeWorkspaceMode::Sandbox); + } + #[cfg(feature = "local-workspace")] + #[test] + fn welcome_ack_cancel_clears_history_bypass() { + use crate::views::welcome::WelcomeWorkspaceMode; + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.auth_state = AuthState::Done; + app.trust_state = TrustState::Done; + app.welcome_local_workspace_ack_pending = true; + app.welcome_workspace_mode = WelcomeWorkspaceMode::LocalWorkspace; + app.welcome_history_load_as_build = true; + app.deferred_startup.worktree = true; + app.deferred_startup.history_load_as_build = true; + let outcome = app.handle_input(&key_event(KeyCode::Char('n'), KeyModifiers::NONE)); + assert!(matches!(outcome, InputOutcome::Changed)); + assert!(!app.welcome_local_workspace_ack_pending); + assert_eq!(app.welcome_workspace_mode, WelcomeWorkspaceMode::Sandbox); + assert!( + !app.welcome_history_load_as_build, + "ACK cancel must drop history bypass" + ); + assert!(!app.deferred_startup.history_load_as_build); + assert!(!app.deferred_startup.worktree); + } + #[cfg(feature = "local-workspace")] + #[test] + fn welcome_workspace_click_selects_mode() { + use crate::views::welcome::{WelcomeWorkspaceMode, WorkspaceModeHitRects}; + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.auth_state = AuthState::Done; + app.trust_state = TrustState::Done; + app.welcome_workspace_mode_rects = WorkspaceModeHitRects { + options: [ + Some(ratatui::layout::Rect::new(10, 5, 9, 1)), + Some(ratatui::layout::Rect::new(20, 5, 17, 1)), + ], + row: Some(ratatui::layout::Rect::new(0, 5, 80, 1)), + }; + let click = Event::Mouse(crossterm::event::MouseEvent { + kind: crossterm::event::MouseEventKind::Down(crossterm::event::MouseButton::Left), + column: 25, + row: 5, + modifiers: KeyModifiers::NONE, + }); + let outcome = app.handle_input(&click); + assert!(matches!(outcome, InputOutcome::Changed)); + assert_eq!( + app.welcome_workspace_mode, + WelcomeWorkspaceMode::LocalWorkspace + ); + } + #[cfg(feature = "local-workspace")] + #[test] + fn welcome_workspace_locked_ignores_cycle_and_click() { + use crate::views::welcome::{WelcomeWorkspaceMode, WorkspaceModeHitRects}; + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.auth_state = AuthState::Done; + app.trust_state = TrustState::Done; + app.local_workspace_startup_locked = true; + app.welcome_workspace_mode = WelcomeWorkspaceMode::LocalWorkspace; + app.welcome_workspace_mode_rects = WorkspaceModeHitRects { + options: [ + Some(ratatui::layout::Rect::new(10, 5, 9, 1)), + Some(ratatui::layout::Rect::new(20, 5, 17, 1)), + ], + row: Some(ratatui::layout::Rect::new(0, 5, 80, 1)), + }; + let key = Event::Key(KeyEvent::new(KeyCode::Char('e'), KeyModifiers::CONTROL)); + assert!(matches!( + app.handle_input(&key), + InputOutcome::Unchanged | InputOutcome::Changed + )); + assert_eq!( + app.welcome_workspace_mode, + WelcomeWorkspaceMode::LocalWorkspace + ); + let click = Event::Mouse(crossterm::event::MouseEvent { + kind: crossterm::event::MouseEventKind::Down(crossterm::event::MouseButton::Left), + column: 12, + row: 5, + modifiers: KeyModifiers::NONE, + }); + let _ = app.handle_input(&click); + assert_eq!( + app.welcome_workspace_mode, + WelcomeWorkspaceMode::LocalWorkspace, + "locked picker must not change selection" + ); + } + #[cfg(feature = "local-workspace")] + #[test] + fn welcome_ctrl_e_ignored_while_history_picker_open() { + use crate::views::welcome::WelcomeWorkspaceMode; + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.auth_state = AuthState::Done; + app.trust_state = TrustState::Done; + app.session_picker_entries = Some(vec![]); + app.session_picker_state.set_query("keep-me"); + let before = app.welcome_workspace_mode; + let key = Event::Key(KeyEvent::new(KeyCode::Char('e'), KeyModifiers::CONTROL)); + let outcome = app.handle_input(&key); + assert_eq!(app.welcome_workspace_mode, before); + assert!( + !matches!(outcome, InputOutcome::Action(Action::ForceDeepSearch)), + "history open: Ctrl+E must not cycle or soft-refresh: {outcome:?}" + ); + assert_eq!(app.session_picker_state.query(), "keep-me"); + let _ = WelcomeWorkspaceMode::Sandbox; + } + #[cfg(feature = "local-workspace")] + #[test] + fn welcome_ctrl_e_ignored_while_authenticating() { + use crate::views::welcome::WelcomeWorkspaceMode; + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.auth_state = AuthState::Authenticating { + request_seq: 1, + handle: None, + auth_url: None, + mode: AuthMode::Command, + }; + app.trust_state = TrustState::Done; + assert_eq!(app.welcome_workspace_mode, WelcomeWorkspaceMode::Sandbox); + let key = Event::Key(KeyEvent::new(KeyCode::Char('e'), KeyModifiers::CONTROL)); + let _ = app.handle_input(&key); + assert_eq!( + app.welcome_workspace_mode, + WelcomeWorkspaceMode::Sandbox, + "Ctrl+E must not cycle mode before auth is Done" + ); + } + #[cfg(feature = "local-workspace")] + #[test] + fn welcome_ctrl_e_ignored_when_zdr_blocked() { + use crate::views::welcome::WelcomeWorkspaceMode; + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.auth_state = AuthState::Done; + app.trust_state = TrustState::Done; + app.is_zdr = true; + app.zdr_access_enabled = false; + assert_eq!(app.welcome_workspace_mode, WelcomeWorkspaceMode::Sandbox); + let key = Event::Key(KeyEvent::new(KeyCode::Char('e'), KeyModifiers::CONTROL)); + let _ = app.handle_input(&key); + assert_eq!( + app.welcome_workspace_mode, + WelcomeWorkspaceMode::Sandbox, + "Ctrl+E must not cycle mode on ZDR-blocked welcome" + ); + } } diff --git a/crates/codegen/xai-grok-pager/src/app/cli.rs b/crates/codegen/xai-grok-pager/src/app/cli.rs index 5e2cee2..2f32d81 100644 --- a/crates/codegen/xai-grok-pager/src/app/cli.rs +++ b/crates/codegen/xai-grok-pager/src/app/cli.rs @@ -606,6 +606,32 @@ pub struct PagerArgs { /// Disable plan mode. #[arg(long = "no-plan")] pub no_plan: bool, + /// Own a local `workspace_server` (replaces remote sandbox). Requires `--chat`. + /// + /// Compiled only with `--features local-workspace` (not implied by `chat`). + #[cfg(feature = "local-workspace")] + #[arg( + long = "local-workspace", + num_args = 0..= 1, + value_name = "CWD", + conflicts_with = "local_workspace_attach", + requires = "chat" + )] + pub local_workspace: Option>, + /// Attach an existing local `workspace_server` by `server_id`, + /// replacing the chat sandbox (ExistingWorkspace only). Requires `--chat`. + #[cfg(feature = "local-workspace")] + #[arg( + long = "local-workspace-attach", + value_name = "SERVER_ID", + conflicts_with = "local_workspace", + requires = "chat" + )] + pub local_workspace_attach: Option, + /// Cwd override for local-workspace attach/own. Requires `--chat`. + #[cfg(feature = "local-workspace")] + #[arg(long = "local-workspace-cwd", value_name = "PATH", requires = "chat")] + pub local_workspace_cwd: Option, /// Disable subagent spawning. #[arg(long = "no-subagents")] pub no_subagents: bool, @@ -824,6 +850,21 @@ impl PagerArgs { pub fn chat(&self) -> bool { false } + /// `--local-workspace[=cwd]` own-mode flag. + #[cfg(feature = "local-workspace")] + pub fn local_workspace(&self) -> Option> { + self.local_workspace.as_ref().map(|inner| inner.as_deref()) + } + /// `--local-workspace-attach=`. + #[cfg(feature = "local-workspace")] + pub fn local_workspace_attach(&self) -> Option<&str> { + self.local_workspace_attach.as_deref() + } + /// `--local-workspace-cwd=`. + #[cfg(feature = "local-workspace")] + pub fn local_workspace_cwd(&self) -> Option<&std::path::Path> { + self.local_workspace_cwd.as_deref() + } /// Get the session ID to resume, from either --resume or --load (hidden alias). /// /// Returns `None` when `--resume` was used without a value (the empty-string diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/ctx.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/ctx.rs index 63bf808..57f8e02 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/ctx.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/ctx.rs @@ -136,6 +136,16 @@ pub(super) fn reseed_tip_for_new_session(app: &mut AppView) { pub(super) fn show_welcome(app: &mut AppView) { app.active_view = ActiveView::Welcome; app.welcome_announcement = WelcomeAnnouncementState::default(); + // Drop stale welcome workspace one-shot / ACK so a later create/load + // cannot inherit an override from a deferred or abandoned NewSession. + #[cfg(feature = "local-workspace")] + { + app.welcome_session_local_workspace = None; + app.welcome_local_workspace_ack_pending = false; + // Abandoning a session/restore must not leak history bypass into the + // next welcome LoadSession / SessionFlags.chat_mode batch. + app.welcome_history_load_as_build = false; + } } /// Restore the view a mid-session auth flow launched from, falling back to the diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/dashboard.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/dashboard.rs index a8445cd..213465b 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/dashboard.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/dashboard.rs @@ -499,7 +499,7 @@ fn clear_pending_overlay_stop(app: &mut AppView) { /// session. /// - First press, any other state (idle, command in flight, cancel /// pending) → arms `AppView::pending_action` with the dashboard's -/// 2s `STOP_CONFIRM_WINDOW`; the shortcuts bar paints "press Ctrl+x +/// 2s `CONFIRM_WINDOW`; the shortcuts bar paints "press Ctrl+x /// again to close this session". Cancel can't help in the non-idle /// variants of this arm — `dispatch_cancel_turn` no-ops unless a /// turn is running, and command cancellation isn't implemented (see @@ -522,7 +522,7 @@ pub(super) fn dispatch_dashboard_overlay_stop(app: &mut AppView) -> Vec if app .agents .get(&id) - .is_some_and(|a| a.session.state.is_turn_running()) + .is_some_and(|a| a.session.state.is_turn_running() || a.session.state.is_compact_running()) { return dispatch_cancel_turn(app); } @@ -1865,7 +1865,7 @@ pub(super) fn dispatch_dashboard_commit_rename(app: &mut AppView) -> Vec /// Without this, closing the selected agent leaves a stale cursor that /// `reanchor_selection` drops to `None`, and the next ↑/↓ restarts from /// the top of the list — a jarring jump. -fn dashboard_neighbor_row( +pub(super) fn dashboard_neighbor_row( app: &AppView, closed: &crate::views::dashboard::DashboardRowId, ) -> Option { @@ -1911,6 +1911,16 @@ fn dashboard_neighbor_row( }) } +/// Ctrl+X on the selected dashboard row, keyed off the row's `RowState` +/// (the same `allows_delete` the renderer paints `[✗]` with): +/// - Deletable row: first press arms, a second within the window deletes. +/// - Busy top-level row: stop what keeps it busy (running turn, background +/// tasks/monitors/`/loop`s, or queued prompts), never arm (a roster row +/// has no local work to stop, so it just reports it must be stopped first). +/// - Subagent row: kill the subagent. +/// +/// Delete only ever runs on an idle row, so it is never queued alongside +/// a `CancelTurn`. pub(super) fn dispatch_dashboard_stop(app: &mut AppView) -> Vec { use crate::views::dashboard::DashboardRowId; use std::time::Instant; @@ -1921,76 +1931,27 @@ pub(super) fn dispatch_dashboard_stop(app: &mut AppView) -> Vec { match &sel { DashboardRowId::TopLevel(id) => { let id = *id; - let Some(agent) = app.agents.get(&id) else { + let Some(agent) = app.agents.get_mut(&id) else { return vec![]; }; - let now = Instant::now(); - // `t.elapsed()` is the idiomatic Instant API - // for "how long since this Instant". Behaviour identical - // to `now.duration_since(*t)` when `t <= now`, which is - // the only case the dispatcher constructs. - let already_confirming = app - .dashboard - .as_ref() - .and_then(|d| d.stop_confirm.as_ref()) - .is_some_and(|(prev, t)| { - *prev == sel - && t.elapsed() < crate::views::dashboard::state::STOP_CONFIRM_WINDOW - }); - if already_confirming { - // Pick the cursor's next home BEFORE the row vanishes, so - // closing moves the selection down 1 instead of letting it - // go stale (which `reanchor_selection` drops to `None`, - // bouncing the next ↑/↓ back to the top of the list). - let neighbor = dashboard_neighbor_row(app, &sel); + if !crate::views::dashboard::classify_top_level(agent).allows_delete() { + // Busy row: stop what keeps it out of Idle — a running turn, + // background work (bg tasks, monitors, scheduled `/loop`s), + // or queued prompts. Never arms; once the row settles to + // idle, Ctrl+X twice deletes it. + let stopped = stop_top_level_activity(agent); if let Some(d) = app.dashboard.as_mut() { - d.stop_confirm = None; + d.delete_confirm = None; } - // Second press: close the agent. - let effects = dispatch_sessions_confirm_close(app, id); - // Only move the cursor if the close actually happened - // (it's refused for the last remaining session). - if !app.agents.contains_key(&id) - && let Some(d) = app.dashboard.as_mut() - { - match neighbor { - Some(n) => d.focus_row(n), - // No neighbour left — land on the always-present - // `[+ New Agent]` button via the focus helper so the - // "exactly one cursor active" invariant holds (a bare - // `selected = None` would leave no cursor and drop the - // footer into its defensive fallback). - None => d.focus_new_agent_button(), + return match stopped { + Some(effects) => effects, + None => { + app.show_toast("Stop the session before deleting"); + vec![] } - } - return effects; + }; } - // First press: cancel turn if running, plant confirmation. - let mut effects = Vec::new(); - if !agent.session.state.is_idle() - && let Some(sid) = agent.session.session_id.clone() - { - effects.push(Effect::CancelTurn { - session_id: sid, - cancel_subagents: true, - trigger: None, - // Dashboard first-press cancel — no local prompt rewind. - rewind_if_pristine: false, - }); - } - if let Some(d) = app.dashboard.as_mut() { - // The footer's `ShortcutsBar::with_pending` already - // paints the "press Ctrl+X again to close this - // session" prompt in the bottom bar. Surfacing the - // same line via `error_toast` would also bleed it - // into the dispatch input placeholder — two copies - // of the same hint, in two different places, with - // the dispatch one stealing visual weight from the - // user's typing area. The footer hint is the - // canonical surface. - d.stop_confirm = Some((sel, now)); - } - effects + arm_or_delete(app, sel) } DashboardRowId::Subagent { parent, @@ -2014,9 +1975,200 @@ pub(super) fn dispatch_dashboard_stop(app: &mut AppView) -> Vec { .into_iter() .collect() } - // Roster-only rows are hosted elsewhere — this client can't stop - // them. - DashboardRowId::Roster { .. } => vec![], + DashboardRowId::Roster { session_id } => { + let entry = app + .leader_roster + .iter() + .chain(app.dashboard_local_sessions.iter()) + .find(|e| e.session_id == session_id.as_str()); + match entry { + None => { + app.show_toast("Session is no longer in the list"); + vec![] + } + // Chat conversations can't be deleted from here yet, so + // don't arm a confirm that could never succeed. + Some(e) if e.origin.kind == "conversation" => { + app.show_toast("Deleting chat conversations isn't supported yet"); + vec![] + } + // No local turn to cancel, so a busy roster row can't delete. + Some(e) + if !crate::views::dashboard::roster_activity_to_state(e.activity) + .allows_delete() => + { + app.show_toast("Stop the session before deleting"); + vec![] + } + Some(_) => arm_or_delete(app, sel), + } + } + } +} + +/// Stop everything keeping a busy top-level row out of Idle: a running +/// turn, running background tasks/monitors, scheduled `/loop`s, and queued +/// (unsent) prompts. Marks local state optimistically (mirroring the agent +/// view's own kill paths). Returns `Some(effects)` when it stopped +/// something — the effects may be empty if the only thing to stop was the +/// local prompt queue — or `None` when there was nothing stoppable (so the +/// caller can explain why). +fn stop_top_level_activity(agent: &mut crate::app::agent_view::AgentView) -> Option> { + let session_id = agent.session.session_id.clone(); + let mut effects = Vec::new(); + + // Turn / background work need a session id to reach the backend. + if let Some(session_id) = session_id { + if !agent.session.state.is_idle() { + effects.push(Effect::CancelTurn { + session_id: session_id.clone(), + cancel_subagents: true, + trigger: None, + rewind_if_pristine: false, + }); + } + let running: Vec = agent + .session + .bg_tasks + .values() + .filter(|t| t.status == crate::app::agent::BgTaskStatus::Running) + .map(|t| t.task_id.clone()) + .collect(); + for task_id in running { + if let Some(task) = agent.session.bg_tasks.get_mut(&task_id) { + task.pending_kill = true; + task.kill_requested_at = Some(std::time::Instant::now()); + } + effects.push(Effect::KillBgTask { + session_id: session_id.clone(), + task_id, + }); + } + let scheduled: Vec = agent.session.scheduled_tasks.keys().cloned().collect(); + for task_id in scheduled { + agent.session.scheduled_tasks.remove(&task_id); + effects.push(Effect::DeleteScheduledTask { + session_id: session_id.clone(), + task_id, + }); + } + } + + // Queued prompts are local (unsent), so dropping them needs no effect + // and works even before the session exists (a just-dispatched row). + let dropped_queue = !agent.session.pending_prompts.is_empty(); + if dropped_queue { + agent.session.pending_prompts.clear(); + agent.sync_queue_pane(); + } + + (!effects.is_empty() || dropped_queue).then_some(effects) +} + +/// A live arm on `sel` confirms and deletes; otherwise (re)arm. +fn arm_or_delete(app: &mut AppView, sel: crate::views::dashboard::DashboardRowId) -> Vec { + let armed = app + .dashboard + .as_mut() + .and_then(|d| d.armed_delete_row()) + .as_ref() + == Some(&sel); + if armed { + return delete_dashboard_row(app, sel); + } + if let Some(d) = app.dashboard.as_mut() { + d.arm_delete(sel); + } + vec![] +} + +pub(super) fn dispatch_dashboard_delete(app: &mut AppView) -> Vec { + let Some(d) = app.dashboard.as_mut() else { + return vec![]; + }; + // The `y` confirm: only fire on a row whose arm is still live AND is + // still the selected row. Reanchor / gc can drop `selected` without + // going through the focus helpers, which would otherwise let `y` + // delete a row the cursor has left. + let Some(sel) = d.armed_delete_row() else { + return vec![]; + }; + if d.selected.as_ref() != Some(&sel) { + d.delete_confirm = None; + return vec![]; + } + delete_dashboard_row(app, sel) +} + +/// Delete `row`, which the caller has confirmed is idle and armed. Takes +/// `row` as a parameter (not read back off `delete_confirm`) and never +/// cancels a turn or kills a task — delete is a settled-row operation. +fn delete_dashboard_row( + app: &mut AppView, + row: crate::views::dashboard::DashboardRowId, +) -> Vec { + use crate::views::dashboard::DashboardRowId; + + if let Some(d) = app.dashboard.as_mut() { + d.delete_confirm = None; + } + match row { + DashboardRowId::TopLevel(id) => { + let Some(agent) = app.agents.get(&id) else { + return vec![]; + }; + // Defensive re-check via the SAME predicate the renderer and + // arm path use: state can change between arming and confirming + // (a new turn, `/loop`, queued prompt, replay, needs-input, or + // bg task), and delete must never run on a non-settled row. + if !crate::views::dashboard::classify_top_level(agent).allows_delete() { + app.show_toast("Stop the session before deleting"); + return vec![]; + } + let Some(session_id) = agent.session.session_id.clone() else { + app.show_toast("No session history to delete"); + return vec![]; + }; + let cwd = agent.session.cwd.display().to_string(); + app.show_toast("Deleting session\u{2026}"); + vec![Effect::DeleteSession { + source: "current".into(), + session_id: session_id.to_string(), + cwd, + after: crate::app::actions::AfterSessionDelete::Dashboard, + }] + } + DashboardRowId::Subagent { .. } => { + app.show_toast("Subagent rows can't be deleted from the dashboard"); + vec![] + } + DashboardRowId::Roster { session_id } => { + let Some(entry) = app + .leader_roster + .iter() + .chain(app.dashboard_local_sessions.iter()) + .find(|e| e.session_id == session_id) + .cloned() + else { + app.show_toast("Session is no longer in the list"); + return vec![]; + }; + if entry.origin.kind == "conversation" { + app.show_toast("Deleting chat conversations isn't supported yet"); + return vec![]; + } + if !crate::views::dashboard::roster_activity_to_state(entry.activity).allows_delete() { + app.show_toast("Stop the session before deleting"); + return vec![]; + } + app.show_toast("Deleting session\u{2026}"); + vec![Effect::DeleteSession { + source: "local".into(), + session_id, + cwd: entry.cwd, + after: crate::app::actions::AfterSessionDelete::Dashboard, + }] + } } } diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/prompt.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/prompt.rs index f9f78c2..1fbd07c 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/prompt.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/prompt.rs @@ -14,7 +14,7 @@ use super::session::fork::open_project_question; use super::session::lifecycle::skip_picker_and_create_session; use super::voice::{merge_prompt_with_voice_interim, voice_stop_on_submit}; use crate::app::actions::{Action, DoctorFixTarget, Effect}; -use crate::app::agent::{AgentId, AgentState}; +use crate::app::agent::{AgentCommand, AgentId, AgentState}; use crate::app::agent_view::AgentView; use crate::app::app_view::{ActiveView, AppView}; use crate::notifications::{NotificationEvent, NotificationEventKind}; @@ -1614,10 +1614,23 @@ pub(super) fn handle_compact_complete( result: Result<(), String>, ) -> Vec { if let Some(agent) = app.agents.get_mut(&agent_id) { - // Defensive: only process if we're still in CommandRunning state. - // This guards against state machine bugs or future cancellation support. - if !matches!(agent.session.state, AgentState::CommandRunning { .. }) { - tracing::debug!("Ignoring CompactComplete (not in CommandRunning state)"); + // Defensive: only process if we're still in a compact command state. + let was_cancelling = matches!( + agent.session.state, + AgentState::CommandCancelling { + command: AgentCommand::Compact, + } + ); + if !matches!( + agent.session.state, + AgentState::CommandRunning { + command: AgentCommand::Compact, + .. + } | AgentState::CommandCancelling { + command: AgentCommand::Compact, + } + ) { + tracing::debug!("Ignoring CompactComplete (not in compact command state)"); return vec![]; } @@ -1632,6 +1645,11 @@ pub(super) fn handle_compact_complete( }, )); } + Err(err) if was_cancelling || err.contains("compact cancelled") => { + agent.scrollback.push_block(RenderBlock::session_event( + SessionEvent::CompactionCancelled, + )); + } Err(err) => { tracing::error!(agent = ?agent_id, error = %err, "Compaction failed"); agent.scrollback.push_block(RenderBlock::session_event( diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/router.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/router.rs index 65e7c6d..050e1f0 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/router.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/router.rs @@ -11,16 +11,17 @@ use super::ctx::{ use super::dashboard::{ dispatch_dashboard_attach, dispatch_dashboard_begin_rename, dispatch_dashboard_change_location, dispatch_dashboard_commit_rename, dispatch_dashboard_confirm_worktree, - dispatch_dashboard_create_new_agent_with_detail, dispatch_dashboard_dispatch, - dispatch_dashboard_dispatch_slash, dispatch_dashboard_open_location_picker, - dispatch_dashboard_open_shortcuts_help, dispatch_dashboard_overlay_cycle, - dispatch_dashboard_overlay_exit, dispatch_dashboard_overlay_stop, - dispatch_dashboard_peek_cycle_mode, dispatch_dashboard_peek_reply, - dispatch_dashboard_permission_followup, dispatch_dashboard_permission_select, - dispatch_dashboard_question_answer, dispatch_dashboard_reorder, dispatch_dashboard_select, - dispatch_dashboard_stop, dispatch_dashboard_toggle_auto_approve, - dispatch_dashboard_toggle_grouping, dispatch_dashboard_toggle_pin, - dispatch_dashboard_toggle_worktree, dispatch_exit_dashboard, dispatch_open_dashboard, + dispatch_dashboard_create_new_agent_with_detail, dispatch_dashboard_delete, + dispatch_dashboard_dispatch, dispatch_dashboard_dispatch_slash, + dispatch_dashboard_open_location_picker, dispatch_dashboard_open_shortcuts_help, + dispatch_dashboard_overlay_cycle, dispatch_dashboard_overlay_exit, + dispatch_dashboard_overlay_stop, dispatch_dashboard_peek_cycle_mode, + dispatch_dashboard_peek_reply, dispatch_dashboard_permission_followup, + dispatch_dashboard_permission_select, dispatch_dashboard_question_answer, + dispatch_dashboard_reorder, dispatch_dashboard_select, dispatch_dashboard_stop, + dispatch_dashboard_toggle_auto_approve, dispatch_dashboard_toggle_grouping, + dispatch_dashboard_toggle_pin, dispatch_dashboard_toggle_worktree, dispatch_exit_dashboard, + dispatch_open_dashboard, }; use super::import_claude::{ dispatch_dismiss_claude_import, dispatch_import_claude, dispatch_import_claude_cancel, @@ -193,6 +194,55 @@ pub(crate) fn dispatch(action: Action, app: &mut AppView) -> Vec { effects } Action::NewSession => dispatch_new_session(app), + #[cfg(feature = "local-workspace")] + Action::ConfirmWelcomeLocalWorkspaceAck => { + match crate::views::welcome::workspace_mode::confirm_welcome_local_workspace_ack( + &app.cwd, false, + ) { + Ok(cfg) => { + app.welcome_workspace_mode = + crate::views::welcome::WelcomeWorkspaceMode::LocalWorkspace; + app.welcome_session_local_workspace = Some(Some(cfg)); + app.welcome_local_workspace_ack_pending = false; + let effects = if app.deferred_startup.worktree { + app.deferred_startup.worktree = false; + let label = app.deferred_startup.worktree_label.take(); + let git_ref = app.deferred_startup.worktree_ref.take(); + let load_session_id = match app.deferred_startup.session.take() { + Some(crate::app::session_startup::DeferredSessionStartup::Load { + session_id, + .. + }) => Some(session_id), + other => { + app.deferred_startup.session = other; + None + } + }; + let preferred = app.deferred_startup.preferred_session_id.take(); + dispatch_new_worktree_session( + app, + load_session_id, + label, + None, + None, + git_ref, + preferred, + ) + } else { + dispatch_new_session(app) + }; + if !crate::app::event_loop::welcome_oneshot_applies_to_effects(&effects) { + app.welcome_session_local_workspace = None; + } + effects + } + Err(err) => { + tracing::warn!("welcome local-workspace ack: {err}"); + app.show_toast(&format!("Local workspace: {err}")); + vec![] + } + } + } Action::ChooseNewSessionMode => open_new_session_question(app), Action::ExitSession | Action::ExitSessionConfirmed => dispatch_exit_session(app), Action::DeleteCurrentSession => open_delete_current_session_question(app), @@ -1260,6 +1310,7 @@ pub(crate) fn dispatch(action: Action, app: &mut AppView) -> Vec { vec![] } Action::DashboardStop => dispatch_dashboard_stop(app), + Action::DashboardDelete => dispatch_dashboard_delete(app), Action::DashboardCycleMode => { let policy_block = app.yolo_policy_block; if let Some(d) = app.dashboard.as_mut() { diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/session/foreign.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/session/foreign.rs index 79aa0dd..2195b8d 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/session/foreign.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/session/foreign.rs @@ -138,6 +138,23 @@ impl PickerSurface<'_> { } } +/// Kind facet for welcome multi-source history under `--chat`. +/// +/// Sandbox → `chat` (gateway); Local → `build` (local-disk). Modal / non-welcome +/// fetches leave this `None` so the shell keeps its default chat-mode force. +pub(in crate::app::dispatch) fn welcome_history_kind_filter(app: &AppView) -> Option> { + #[cfg(feature = "local-workspace")] + { + if app.chat_mode && matches!(app.active_view, crate::app::app_view::ActiveView::Welcome) { + return Some(vec![ + app.welcome_workspace_mode.history_kind_filter().to_string(), + ]); + } + } + let _ = app; + None +} + pub(in crate::app::dispatch) fn dispatch_fetch_session_list(app: &mut AppView) -> Vec { app.session_picker_detail_generation += 1; app.session_picker_loading = true; @@ -154,9 +171,18 @@ pub(in crate::app::dispatch) fn dispatch_fetch_session_list(app: &mut AppView) - } app.foreign_session_scan_seq += 1; let foreign_seq = app.foreign_session_scan_seq; + let kind_filter = welcome_history_kind_filter(app); + #[cfg(feature = "local-workspace")] + crate::views::welcome::workspace_mode::log_history_source( + "session_list_fetch_dispatch", + Some(app.welcome_workspace_mode), + kind_filter.as_deref(), + None, + ); let mut effects = vec![Effect::FetchSessionList { query: None, seq: app.session_picker_list_seq, + kind_filter, }]; let foreign_effect = if app.chat_mode { app.foreign_scan_coordinator.begin_request(foreign_seq); @@ -269,20 +295,19 @@ pub(in crate::app::dispatch) fn handle_session_list_loaded( app.show_toast(¬ice); } else if scope.is_relaxed() && app.session_picker_relaxed_notified_for.as_deref() != Some(app.cwd.as_path()) - { // Welcome view drops toasts; don't consume the one-shot notice unless // it can render. - if !matches!(app.active_view, crate::app::app_view::ActiveView::Welcome) { - // Notify once per directory; the browse is scoped to `app.cwd`. - app.session_picker_relaxed_notified_for = Some(app.cwd.clone()); - let message = match scope { - ListScope::Repo => { - "No sessions in this directory. Showing other sessions from this repository." - } - _ => "No sessions in this directory. Showing sessions from other directories.", - }; - app.show_toast(message); - } + && !matches!(app.active_view, crate::app::app_view::ActiveView::Welcome) + { + // Notify once per directory; the browse is scoped to `app.cwd`. + app.session_picker_relaxed_notified_for = Some(app.cwd.clone()); + let message = match scope { + ListScope::Repo => { + "No sessions in this directory. Showing other sessions from this repository." + } + _ => "No sessions in this directory. Showing sessions from other directories.", + }; + app.show_toast(message); } // A cwd-scoped browse clears the latch so a later relax re-notifies; search // responses leave it alone. diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/session/lifecycle.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/session/lifecycle.rs index 9fb97ec..8a61393 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/session/lifecycle.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/session/lifecycle.rs @@ -127,6 +127,13 @@ pub(in crate::app::dispatch) fn dispatch_new_session(app: &mut AppView) -> Vec Result<(), Vec> { + use crate::views::welcome::workspace_mode::{ + WelcomeWorkspaceMode, WelcomeWorkspacePrepare, prepare_welcome_workspace_for_new_session, + }; + match prepare_welcome_workspace_for_new_session( + app.welcome_workspace_mode, + app.local_workspace_startup_locked, + app.chat_mode, + &app.cwd, + false, + ) { + Ok(WelcomeWorkspacePrepare::Continue { + session_override, + warning, + }) => { + if let Some(msg) = warning { + tracing::warn!("{msg}"); + app.show_toast(&msg); + } + if let Some(override_cfg) = session_override { + app.welcome_session_local_workspace = Some(override_cfg); + } + Ok(()) + } + Ok(WelcomeWorkspacePrepare::AwaitAck) => { + app.welcome_local_workspace_ack_pending = true; + app.session_picker_entries = None; + app.session_picker_loading = false; + app.session_picker_list_seq = app.session_picker_list_seq.saturating_add(1); + Err(vec![]) + } + Err(err) => { + tracing::warn!("welcome workspace mode: {err}"); + app.show_toast(&format!( + "Local workspace unavailable ({err}); using sandbox" + )); + app.welcome_session_local_workspace = Some(None); + app.welcome_workspace_mode = WelcomeWorkspaceMode::Sandbox; + Ok(()) + } + } +} /// Factored out of [`dispatch_new_session`] so the worktree-question /// "No" path can call it directly without re-opening the modal. pub(in crate::app::dispatch) fn dispatch_new_session_inner( @@ -362,6 +416,26 @@ 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; + #[cfg(feature = "local-workspace")] + { + let local_intent = match &app.welcome_session_local_workspace { + Some(Some(_)) => true, + Some(None) => false, + None => crate::app::session_startup::active_local_workspace() + .ok() + .flatten() + .is_some(), + }; + let (mode, locked) = + crate::views::welcome::workspace_mode::indicator_for_opening_session( + agent.chat_kind, + false, + app.local_workspace_startup_locked, + local_intent, + ); + agent.workspace_mode = mode; + agent.workspace_mode_cli_locked = locked; + } agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); agent.mcp_init_progress = Some(McpInitProgress { total: 0, @@ -549,6 +623,8 @@ pub(in crate::app::dispatch) fn drain_startup_actions(app: &mut AppView) -> Vec< prompt, open_dashboard, pending_chat, + #[cfg(feature = "local-workspace")] + history_load_as_build, } = app.deferred_startup.take(); let mut effects = Vec::new(); match deferred { @@ -569,6 +645,10 @@ pub(in crate::app::dispatch) fn drain_startup_actions(app: &mut AppView) -> Vec< session_cwd, chat_kind, }) => { + #[cfg(feature = "local-workspace")] + { + app.welcome_history_load_as_build = history_load_as_build; + } if worktree { if chat_kind || pending_chat { app.deferred_startup.pending_chat = true; @@ -672,6 +752,33 @@ pub(in crate::app::dispatch) fn dispatch_new_worktree_session( git_ref: Option, preferred_session_id: Option, ) -> Vec { + #[cfg(feature = "local-workspace")] + if load_session_id.is_none() + && matches!(app.active_view, crate::app::app_view::ActiveView::Welcome) + { + let skip_apply = app.welcome_session_local_workspace.is_some(); + if !skip_apply && let Err(effects) = apply_welcome_workspace_on_new_session(app) { + app.deferred_startup.worktree = true; + if let Some(ref label) = label { + app.deferred_startup.worktree_label = Some(label.clone()); + } + if let Some(ref git_ref) = git_ref { + app.deferred_startup.worktree_ref = Some(git_ref.clone()); + } + if let Some(sid) = load_session_id.clone() { + app.deferred_startup.session = + Some(crate::app::session_startup::DeferredSessionStartup::Load { + session_id: sid, + session_cwd: None, + chat_kind: app.deferred_startup.pending_chat, + }); + } + if let Some(id) = preferred_session_id.clone() { + app.deferred_startup.preferred_session_id = Some(id); + } + return effects; + } + } let preferred_session_id = preferred_session_id.or_else(|| app.deferred_startup.preferred_session_id.take()); if !app.session_startup_allowed() { @@ -710,6 +817,11 @@ pub(in crate::app::dispatch) fn dispatch_new_worktree_session( action: None, }); } + #[cfg(feature = "local-workspace")] + { + app.welcome_session_local_workspace = None; + app.welcome_history_load_as_build = false; + } return vec![]; } if load_session_id.is_none() { @@ -791,6 +903,26 @@ pub(in crate::app::dispatch) fn dispatch_new_worktree_session( &app.tier_restricted_commands, ); agent.chat_kind = chat_kind; + #[cfg(feature = "local-workspace")] + { + let local_intent = match &app.welcome_session_local_workspace { + Some(Some(_)) => true, + Some(None) => false, + None => crate::app::session_startup::active_local_workspace() + .ok() + .flatten() + .is_some(), + }; + let (mode, locked) = + crate::views::welcome::workspace_mode::indicator_for_opening_session( + agent.chat_kind, + app.welcome_history_load_as_build, + app.local_workspace_startup_locked, + local_intent, + ); + agent.workspace_mode = mode; + agent.workspace_mode_cli_locked = locked; + } agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); agent .prompt @@ -880,6 +1012,26 @@ 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; + #[cfg(feature = "local-workspace")] + { + let local_intent = match &app.welcome_session_local_workspace { + Some(Some(_)) => true, + Some(None) => false, + None => crate::app::session_startup::active_local_workspace() + .ok() + .flatten() + .is_some(), + }; + let (mode, locked) = + crate::views::welcome::workspace_mode::indicator_for_opening_session( + agent.chat_kind, + false, + app.local_workspace_startup_locked, + local_intent, + ); + agent.workspace_mode = mode; + agent.workspace_mode_cli_locked = locked; + } agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); agent.mcp_init_progress = Some(McpInitProgress { total: 0, diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/session/load.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/session/load.rs index f3cd940..3138ff8 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/session/load.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/session/load.rs @@ -9,6 +9,8 @@ use crate::acp::tracker::AcpUpdateTracker; use crate::app::actions::{Action, Effect}; use crate::app::agent::{AgentCommand, AgentId, AgentSession, AgentState}; use crate::app::agent_view::AgentView; +#[cfg(feature = "local-workspace")] +use crate::app::app_view::ActiveView; use crate::app::app_view::AppView; use crate::app::dispatch::ctx::{ SwitchCause, get_active_agent, get_active_agent_mut, switch_to_agent, with_active_agent, @@ -34,6 +36,11 @@ pub(in crate::app::dispatch) fn dispatch_load_session( chat_kind: bool, ) -> Vec { if !app.session_startup_allowed() { + #[cfg(feature = "local-workspace")] + { + app.deferred_startup.history_load_as_build = app.welcome_history_load_as_build; + app.welcome_history_load_as_build = false; + } app.deferred_startup.session = Some(crate::app::session_startup::DeferredSessionStartup::Load { session_id, @@ -117,17 +124,31 @@ fn dispatch_load_session_ungated( session_cwd: Option, chat_kind: bool, ) -> Vec { - if crate::app::session_startup::chat_mode_refuses_local_build_load( - app.chat_mode, - chat_kind, - &session_id, - &app.cwd, - ) { + #[cfg(feature = "local-workspace")] + let bypass_chat_refusal = app.welcome_history_load_as_build; + #[cfg(not(feature = "local-workspace"))] + let bypass_chat_refusal = false; + if !bypass_chat_refusal + && crate::app::session_startup::chat_mode_refuses_local_build_load( + app.chat_mode, + chat_kind, + &session_id, + &app.cwd, + ) + { + #[cfg(feature = "local-workspace")] + { + app.welcome_history_load_as_build = false; + } app.show_toast(crate::app::session_startup::CHAT_MODE_LOCAL_BUILD_REFUSAL); return vec![]; } invalidate_picker_fetch_on_dismiss(app); if focus_if_session_already_open(app, &session_id, chat_kind).is_some() { + #[cfg(feature = "local-workspace")] + { + app.welcome_history_load_as_build = false; + } return vec![]; } let acp_session_id = clear_stale_session_id(app, &session_id); @@ -210,6 +231,33 @@ fn dispatch_load_session_ungated( &app.tier_restricted_commands, ); agent_mut.chat_kind = chat_kind || app.chat_mode; + #[cfg(feature = "local-workspace")] + { + let history_build = app.welcome_history_load_as_build; + let local_intent = match &app.welcome_session_local_workspace { + Some(Some(_)) => true, + Some(None) => false, + None => { + if chat_kind { + false + } else { + crate::app::session_startup::active_local_workspace() + .ok() + .flatten() + .is_some() + } + } + }; + let (mode, cli_locked) = + crate::views::welcome::workspace_mode::indicator_for_opening_session( + chat_kind, + history_build, + app.local_workspace_startup_locked, + local_intent, + ); + agent_mut.workspace_mode = mode; + agent_mut.workspace_mode_cli_locked = cli_locked; + } agent_mut.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); agent_mut .prompt @@ -313,6 +361,29 @@ pub(in crate::app::dispatch) fn dispatch_pick_session( return effects; } let chat_kind = source == "conversation"; + #[cfg(feature = "local-workspace")] + if app.chat_mode && matches!(app.active_view, ActiveView::Welcome) { + if app.local_workspace_startup_locked { + crate::views::welcome::workspace_mode::log_cli_lock_wins(app.welcome_workspace_mode); + } else { + let mode = crate::views::welcome::WelcomeWorkspaceMode::from_history_source(&source); + if app.welcome_workspace_mode != mode { + crate::views::welcome::workspace_mode::log_history_source( + "history_auto_switch", + Some(mode), + None, + Some(source.as_str()), + ); + app.welcome_workspace_mode = mode; + } + if chat_kind { + app.welcome_session_local_workspace = None; + } + } + if !chat_kind { + app.welcome_history_load_as_build = true; + } + } if chat_kind { return dispatch_load_session(app, session_id, None, true); } @@ -331,11 +402,19 @@ pub(in crate::app::dispatch) fn dispatch_pick_session( } if source == "remote" || source == "both" { if focus_if_session_already_open(app, &session_id, false).is_some() { + #[cfg(feature = "local-workspace")] + { + app.welcome_history_load_as_build = false; + } return vec![]; } app.show_toast("Restoring session from remote..."); dispatch_load_session_with_restore(app, session_id, cwd) } else { + #[cfg(feature = "local-workspace")] + { + app.welcome_history_load_as_build = false; + } app.show_toast("Session not found locally"); vec![] } @@ -412,6 +491,24 @@ pub(in crate::app::dispatch) fn dispatch_pick_session_in_worktree( app.show_toast("Chat conversations can't be resumed in a worktree"); return vec![]; } + #[cfg(feature = "local-workspace")] + if app.chat_mode && matches!(app.active_view, ActiveView::Welcome) { + if app.local_workspace_startup_locked { + crate::views::welcome::workspace_mode::log_cli_lock_wins(app.welcome_workspace_mode); + } else { + let mode = crate::views::welcome::WelcomeWorkspaceMode::from_history_source(&source); + if app.welcome_workspace_mode != mode { + crate::views::welcome::workspace_mode::log_history_source( + "history_auto_switch", + Some(mode), + None, + Some(source.as_str()), + ); + app.welcome_workspace_mode = mode; + } + } + app.welcome_history_load_as_build = true; + } dispatch_new_worktree_session(app, Some(session_id), None, None, None, None, None) } fn keep_picker_entry( @@ -454,9 +551,7 @@ pub(in crate::app::dispatch) fn remove_session_from_pickers( { if pending_delete .as_ref() - .is_some_and(|(pending_source, pending_id, _)| { - pending_source == source && pending_id == session_id - }) + .is_some_and(|pd| pd.source == source && pd.session_id == session_id) { *pending_delete = None; } @@ -482,6 +577,13 @@ pub(in crate::app::dispatch) fn remove_session_from_pickers( ); reanchor_grouped_selection(state, &map); } + if app + .session_picker_pending_delete + .as_ref() + .is_some_and(|pd| pd.source == source && pd.session_id == session_id) + { + app.session_picker_pending_delete = None; + } if let Some(list) = app.session_picker_entries.as_mut() { list.retain(|entry| keep_picker_entry(entry, source, session_id, match_id_only)); } @@ -653,13 +755,18 @@ fn dispatch_chat_search_refetch(app: &mut AppView, force: bool) -> Vec { let seq = app.session_picker_list_seq; if query.is_empty() { set_chat_search_loading(app, false); - return vec![Effect::FetchSessionList { query: None, seq }]; + return vec![Effect::FetchSessionList { + query: None, + seq, + kind_filter: super::foreign::welcome_history_kind_filter(app), + }]; } set_chat_search_loading(app, true); if force { vec![Effect::FetchSessionList { query: Some(query), seq, + kind_filter: super::foreign::welcome_history_kind_filter(app), }] } else { vec![Effect::DebounceSessionSearch { query, seq }] @@ -789,16 +896,30 @@ pub(in crate::app::dispatch) fn dispatch_load_session_with_restore( session_id: String, session_cwd: String, ) -> Vec { - if crate::app::session_startup::chat_mode_refuses_local_build_load( - app.chat_mode, - false, - &session_id, - &app.cwd, - ) { + #[cfg(feature = "local-workspace")] + let bypass_chat_refusal = app.welcome_history_load_as_build; + #[cfg(not(feature = "local-workspace"))] + let bypass_chat_refusal = false; + if !bypass_chat_refusal + && crate::app::session_startup::chat_mode_refuses_local_build_load( + app.chat_mode, + false, + &session_id, + &app.cwd, + ) + { + #[cfg(feature = "local-workspace")] + { + app.welcome_history_load_as_build = false; + } app.show_toast(crate::app::session_startup::CHAT_MODE_LOCAL_BUILD_REFUSAL); return vec![]; } if focus_if_session_already_open(app, &session_id, false).is_some() { + #[cfg(feature = "local-workspace")] + { + app.welcome_history_load_as_build = false; + } return vec![]; } let agent_id = AgentId(app.next_agent_id); @@ -870,6 +991,27 @@ pub(in crate::app::dispatch) fn dispatch_load_session_with_restore( &app.tier_restricted_commands, ); agent.chat_kind = app.chat_mode; + #[cfg(feature = "local-workspace")] + { + let history_build = app.welcome_history_load_as_build; + let local_intent = match &app.welcome_session_local_workspace { + Some(Some(_)) => true, + Some(None) => false, + None => crate::app::session_startup::active_local_workspace() + .ok() + .flatten() + .is_some(), + }; + let (mode, cli_locked) = + crate::views::welcome::workspace_mode::indicator_for_opening_session( + false, + history_build, + app.local_workspace_startup_locked, + local_intent, + ); + agent.workspace_mode = mode; + agent.workspace_mode_cli_locked = cli_locked; + } agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); agent .prompt @@ -1068,6 +1210,7 @@ pub(in crate::app::dispatch) fn handle_session_search_debounce_expired( return vec![Effect::FetchSessionList { query: (!query.is_empty()).then_some(query), seq, + kind_filter: super::foreign::welcome_history_kind_filter(app), }]; } if live_deep_search_seq(app) != Some(seq) { @@ -1136,12 +1279,22 @@ pub(in crate::app::dispatch) fn handle_session_restored( agent_id: AgentId, local_session_id: String, ) -> Vec { - if crate::app::session_startup::chat_mode_refuses_local_build_load( - app.chat_mode, - false, - &local_session_id, - &app.cwd, - ) { + #[cfg(feature = "local-workspace")] + let bypass_chat_refusal = app.welcome_history_load_as_build; + #[cfg(not(feature = "local-workspace"))] + let bypass_chat_refusal = false; + if !bypass_chat_refusal + && crate::app::session_startup::chat_mode_refuses_local_build_load( + app.chat_mode, + false, + &local_session_id, + &app.cwd, + ) + { + #[cfg(feature = "local-workspace")] + { + app.welcome_history_load_as_build = false; + } refuse_chat_mode_build_agent(app, agent_id); return vec![]; } @@ -1150,6 +1303,27 @@ 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; + #[cfg(feature = "local-workspace")] + { + let history_build = app.welcome_history_load_as_build; + let local_intent = match &app.welcome_session_local_workspace { + Some(Some(_)) => true, + Some(None) => false, + None => crate::app::session_startup::active_local_workspace() + .ok() + .flatten() + .is_some(), + }; + let (mode, cli_locked) = + crate::views::welcome::workspace_mode::indicator_for_opening_session( + false, + history_build, + app.local_workspace_startup_locked, + local_intent, + ); + agent.workspace_mode = mode; + agent.workspace_mode_cli_locked = cli_locked; + } 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}..." @@ -1170,6 +1344,10 @@ pub(in crate::app::dispatch) fn handle_session_restore_failed( error: String, ) -> Vec { tracing::error!(agent = ?agent_id, error = %error, "Session restore failed"); + #[cfg(feature = "local-workspace")] + { + app.welcome_history_load_as_build = false; + } if let Some(agent) = app.agents.get_mut(&agent_id) { if defer_to_open_reload_window(agent, agent_id, "SessionRestoreFailed") { return vec![]; diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/task_result.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/task_result.rs index 239c96d..4622508 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/task_result.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/task_result.rs @@ -910,6 +910,10 @@ pub(super) fn dispatch_task_result(result: TaskResult, app: &mut AppView) -> Vec after != AfterSessionDelete::Stay, ); if after == AfterSessionDelete::Stay { + app.dashboard_local_sessions + .retain(|entry| entry.session_id != session_id); + app.leader_roster + .retain(|entry| entry.session_id != session_id); app.show_toast("Session deleted"); return vec![]; } @@ -922,11 +926,49 @@ pub(super) fn dispatch_task_result(result: TaskResult, app: &mut AppView) -> Vec .collect(); let foreground = matches!(app.active_view, ActiveView::Agent(id) if to_remove.contains(&id)); + let roster_row = crate::views::dashboard::DashboardRowId::Roster { + session_id: session_id.clone(), + }; + let closed_rows: Vec<_> = to_remove + .iter() + .copied() + .map(crate::views::dashboard::DashboardRowId::TopLevel) + .chain(std::iter::once(roster_row)) + .collect(); + let selected = app.dashboard.as_ref().and_then(|d| d.selected.clone()); + let neighbor = if after == AfterSessionDelete::Dashboard + && let Some(sel) = selected.as_ref().filter(|sel| closed_rows.contains(sel)) + { + super::dashboard::dashboard_neighbor_row(app, sel) + } else { + None + }; + app.dashboard_local_sessions + .retain(|entry| entry.session_id != session_id); + app.leader_roster + .retain(|entry| entry.session_id != session_id); for id in to_remove { remove_agent_and_cleanup(app, id); } let mut effects = unregister_session_effect(Some(sid)); - if foreground && after == AfterSessionDelete::Welcome { + if after == AfterSessionDelete::Dashboard { + if let Some(d) = app.dashboard.as_mut() { + d.delete_confirm = None; + let selected_closed = d + .selected + .as_ref() + .is_some_and(|sel| closed_rows.contains(sel)); + match (selected_closed, neighbor) { + (true, Some(n)) => d.focus_row(n), + (true, None) => d.focus_new_agent_button(), + _ => {} + } + } + if foreground { + super::dashboard::ensure_dashboard_state(app); + app.active_view = ActiveView::AgentDashboard; + } + } else if foreground && after == AfterSessionDelete::Welcome { effects.extend(dispatch_exit_session(app)); } app.show_toast("Session deleted"); diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/dashboard.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/dashboard.rs index 27a6889..f1f7f32 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/dashboard.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/dashboard.rs @@ -3038,14 +3038,11 @@ fn dashboard_overlay_stop_busy_agent_cancels_instead_of_closing() { "the overlay attachment must survive", ); } -/// A COMMAND in flight at confirm time must NOT downgrade to a -/// cancel — `dispatch_cancel_turn` no-ops for command states, which -/// would silently eat the confirmed press. The close proceeds: it -/// is the only termination the user can reach (commands can't be -/// cancelled). +/// `/compact` in flight: overlay stop cancels compaction instead of +/// closing the session (same as a running turn). #[serial_test::serial(GROK_AGENT_DASHBOARD)] #[test] -fn dashboard_overlay_stop_command_running_closes_session() { +fn dashboard_overlay_stop_compact_running_cancels() { let mut app = test_app_with_agent(); mark_agent_nonempty(&mut app, AgentId(0)); let id2 = AgentId(1); @@ -3062,15 +3059,27 @@ fn dashboard_overlay_stop_command_running_closes_session() { command: crate::app::agent::AgentCommand::Compact, started_at: std::time::Instant::now(), }; - let _ = dispatch_dashboard_overlay_stop(&mut app); + app.agents.get_mut(&id).unwrap().session.session_id = Some(acp::SessionId::new("sess-compact")); + let effects = dispatch_dashboard_overlay_stop(&mut app); assert!( - !app.agents.contains_key(&id), - "a command in flight must not block the confirmed close", + app.agents.contains_key(&id), + "stop during /compact must not close the session", ); assert!( - matches!(app.active_view, ActiveView::AgentDashboard), - "the confirmed close must land on the dashboard, got {:?}", - app.active_view, + matches!( + app.agents.get(&id).unwrap().session.state, + AgentState::CommandCancelling { + command: crate::app::agent::AgentCommand::Compact, + } + ), + "stop during /compact must enter CommandCancelling, got {:?}", + app.agents.get(&id).unwrap().session.state, + ); + assert!( + effects + .iter() + .any(|e| matches!(e, Effect::CancelTurn { .. })), + "stop during /compact must emit CancelTurn, got {effects:?}", ); } /// An armed overlay stop-confirm is bound to "this overlay, this @@ -4010,25 +4019,27 @@ fn dashboard_open_drops_pinned_ids_for_missing_agents() { let d = app.dashboard.as_ref().unwrap(); assert!(d.pinned.is_empty(), "stale pin should be gc'd at open"); } -/// Ctrl+X first press arms confirm, second -/// press within 2s closes. We don't sleep — we manually rewind -/// `stop_confirm.1` to a recent instant and check the second -/// press is honoured. #[serial_test::serial(GROK_AGENT_DASHBOARD)] #[test] -fn dashboard_stop_double_press_closes_top_level() { +fn dashboard_stop_double_press_deletes_top_level() { let mut app = test_app(); let _ = dispatch_new_session_inner(&mut app, None); let _ = dispatch_new_session_inner(&mut app, None); + for (i, a) in app.agents.values_mut().enumerate() { + a.session.session_id = Some(acp::SessionId::new(format!("s{i}"))); + } open_dashboard(&mut app); let target = *app.agents.keys().next().unwrap(); if let Some(d) = app.dashboard.as_mut() { d.selected = Some(crate::views::dashboard::DashboardRowId::TopLevel(target)); } let _ = dispatch_dashboard_stop(&mut app); - assert!(app.dashboard.as_ref().unwrap().stop_confirm.is_some()); - let _ = dispatch_dashboard_stop(&mut app); - assert!(!app.agents.contains_key(&target)); + assert!(app.dashboard.as_ref().unwrap().delete_confirm.is_some()); + let effects = dispatch_dashboard_stop(&mut app); + assert!(matches!( + effects.last(), + Some(crate::app::actions::Effect::DeleteSession { .. }) + )); } /// Closing the selected agent moves the cursor DOWN one row (onto the /// agent that shifts up into its place) instead of dropping it to @@ -4042,6 +4053,7 @@ fn dashboard_stop_moves_selection_down_one() { let _ = dispatch_new_session_inner(&mut app, None); for (i, agent) in app.agents.values_mut().enumerate() { agent.display_name = Some(format!("agent-{i}")); + agent.session.session_id = Some(acp::SessionId::new(format!("s{i}"))); } open_dashboard(&mut app); let order = dashboard_row_order(&app); @@ -4052,19 +4064,33 @@ fn dashboard_stop_moves_selection_down_one() { d.focus_row(first.clone()); } let _ = dispatch_dashboard_stop(&mut app); - let _ = dispatch_dashboard_stop(&mut app); - let crate::views::dashboard::DashboardRowId::TopLevel(first_id) = first else { + let effects = dispatch_dashboard_stop(&mut app); + let crate::views::dashboard::DashboardRowId::TopLevel(first_id) = &first else { panic!("first row should be top-level"); }; + let session_id = app.agents[first_id] + .session + .session_id + .as_ref() + .expect("session id") + .to_string(); assert!( - !app.agents.contains_key(&first_id), - "closed agent must be gone" + matches!( + effects.last(), + Some(crate::app::actions::Effect::DeleteSession { .. }) + ), + "second Ctrl+X must delete, got {effects:?}" ); - assert_eq!( - app.dashboard.as_ref().unwrap().selected, - Some(second), - "closing the top row should select the next row down, not revert to top", + let _ = dispatch_task_result( + crate::app::actions::TaskResult::DeleteSessionComplete { + source: "current".into(), + session_id, + after: crate::app::actions::AfterSessionDelete::Dashboard, + }, + &mut app, ); + assert!(!app.agents.contains_key(first_id)); + assert_eq!(app.dashboard.as_ref().unwrap().selected, Some(second)); } /// Closing the LAST row has no row below it, so the cursor falls back /// to the previous row rather than disappearing. @@ -4077,6 +4103,7 @@ fn dashboard_stop_last_row_falls_back_to_previous() { let _ = dispatch_new_session_inner(&mut app, None); for (i, agent) in app.agents.values_mut().enumerate() { agent.display_name = Some(format!("agent-{i}")); + agent.session.session_id = Some(acp::SessionId::new(format!("s{i}"))); } open_dashboard(&mut app); let order = dashboard_row_order(&app); @@ -4087,25 +4114,36 @@ fn dashboard_stop_last_row_falls_back_to_previous() { d.focus_row(last.clone()); } let _ = dispatch_dashboard_stop(&mut app); - let _ = dispatch_dashboard_stop(&mut app); - let crate::views::dashboard::DashboardRowId::TopLevel(last_id) = last else { + let effects = dispatch_dashboard_stop(&mut app); + let crate::views::dashboard::DashboardRowId::TopLevel(last_id) = &last else { panic!("last row should be top-level"); }; - assert!( - !app.agents.contains_key(&last_id), - "closed agent must be gone" - ); - assert_eq!( - app.dashboard.as_ref().unwrap().selected, - Some(prev), - "closing the last row should select the previous row", + let session_id = app.agents[last_id] + .session + .session_id + .as_ref() + .expect("session id") + .to_string(); + assert!(matches!( + effects.last(), + Some(crate::app::actions::Effect::DeleteSession { .. }) + )); + let _ = dispatch_task_result( + crate::app::actions::TaskResult::DeleteSessionComplete { + source: "current".into(), + session_id, + after: crate::app::actions::AfterSessionDelete::Dashboard, + }, + &mut app, ); + assert!(!app.agents.contains_key(last_id)); + assert_eq!(app.dashboard.as_ref().unwrap().selected, Some(prev)); } /// First Ctrl+X must NOT plant an `error_toast`. The /// dispatch-input placeholder is reserved for the user's typing /// target — the footer's `ShortcutsBar::with_pending` already /// surfaces the "press Ctrl+X again to close this session" -/// hint via `stop_confirm` and is the canonical place for it. +/// hint via `delete_confirm` and is the canonical place for it. /// Two copies of the same hint in two different surfaces /// confused the user (the prompt one stole visual weight). #[serial_test::serial(GROK_AGENT_DASHBOARD)] @@ -4122,8 +4160,8 @@ fn dashboard_stop_does_not_plant_error_toast() { let _ = dispatch_dashboard_stop(&mut app); let d = app.dashboard.as_ref().unwrap(); assert!( - d.stop_confirm.is_some(), - "first Ctrl+X must arm stop_confirm (footer reads from this)", + d.delete_confirm.is_some(), + "first Ctrl+X on an idle row must arm delete_confirm (footer reads from this)", ); assert!( d.error_toast.is_none(), @@ -4488,14 +4526,14 @@ fn dashboard_stop_double_press_after_2s_rearms() { } let _ = dispatch_dashboard_stop(&mut app); if let Some(d) = app.dashboard.as_mut() - && let Some((_row, at)) = d.stop_confirm.as_mut() + && let Some((_row, at)) = d.delete_confirm.as_mut() { *at = Instant::now() - Duration::from_secs(3); } let before_count = app.agents.len(); let _ = dispatch_dashboard_stop(&mut app); assert_eq!(app.agents.len(), before_count); - assert!(app.dashboard.as_ref().unwrap().stop_confirm.is_some()); + assert!(app.dashboard.as_ref().unwrap().delete_confirm.is_some()); } /// Subagent Ctrl+X bypasses confirm and emits KillSubagent. #[serial_test::serial(GROK_AGENT_DASHBOARD)] @@ -4519,7 +4557,248 @@ fn dashboard_stop_subagent_emits_kill_subagent_effect() { effects.as_slice(), [Effect::KillSubagent { subagent_id, .. }] if subagent_id == "sa-xyz" )); - assert!(app.dashboard.as_ref().unwrap().stop_confirm.is_none()); + assert!(app.dashboard.as_ref().unwrap().delete_confirm.is_none()); +} +#[serial_test::serial(GROK_AGENT_DASHBOARD)] +#[test] +fn dashboard_delete_complete_returns_from_foreground_agent() { + let mut app = test_app_with_agent(); + let id = AgentId(0); + app.agents.get_mut(&id).unwrap().session.session_id = Some(acp::SessionId::new("sess-dash")); + open_dashboard(&mut app); + app.active_view = ActiveView::Agent(id); + let _ = dispatch_task_result( + crate::app::actions::TaskResult::DeleteSessionComplete { + source: "current".into(), + session_id: "sess-dash".into(), + after: crate::app::actions::AfterSessionDelete::Dashboard, + }, + &mut app, + ); + assert!(!app.agents.contains_key(&id)); + assert!(matches!(app.active_view, ActiveView::AgentDashboard)); +} +#[serial_test::serial(GROK_AGENT_DASHBOARD)] +#[test] +fn dashboard_stop_busy_roster_toasts_without_arming() { + let mut app = test_app(); + let mut entry = idle_roster_entry("busy-sess", "Busy row"); + entry.activity = crate::app::roster::RosterActivity::Working; + app.dashboard_local_sessions = vec![entry]; + open_dashboard(&mut app); + if let Some(d) = app.dashboard.as_mut() { + d.selected = Some(crate::views::dashboard::DashboardRowId::Roster { + session_id: "busy-sess".into(), + }); + } + assert!(dispatch_dashboard_stop(&mut app).is_empty()); + let d = app.dashboard.as_ref().unwrap(); + assert!(d.delete_confirm.is_none()); + assert_eq!( + d.error_toast.as_deref(), + Some("Stop the session before deleting"), + ); +} +/// A busy top-level row: Ctrl+X cancels the turn and never arms delete. +#[serial_test::serial(GROK_AGENT_DASHBOARD)] +#[test] +fn dashboard_stop_busy_top_level_cancels_without_arming() { + let mut app = test_app(); + let _ = dispatch_new_session_inner(&mut app, None); + let _ = dispatch_new_session_inner(&mut app, None); + let target = *app.agents.keys().next().unwrap(); + { + let agent = app.agents.get_mut(&target).unwrap(); + agent.session.session_id = Some(acp::SessionId::new("busy-top")); + agent.session.state = AgentState::TurnRunning; + } + open_dashboard(&mut app); + if let Some(d) = app.dashboard.as_mut() { + d.selected = Some(crate::views::dashboard::DashboardRowId::TopLevel(target)); + } + let effects = dispatch_dashboard_stop(&mut app); + assert!( + matches!(effects.as_slice(), [Effect::CancelTurn { .. }]), + "busy top-level Ctrl+X must cancel the turn, got {effects:?}", + ); + assert!( + app.dashboard.as_ref().unwrap().delete_confirm.is_none(), + "busy top-level Ctrl+X must NOT arm delete", + ); + let effects = dispatch_dashboard_stop(&mut app); + assert!( + !effects + .iter() + .any(|e| matches!(e, Effect::DeleteSession { .. })), + "a busy row must never emit DeleteSession, got {effects:?}", + ); + assert!(app.agents.contains_key(&target), "busy row must survive"); +} +/// A row that's `Working` only due to background work (turn idle, a +/// scheduled `/loop` live): Ctrl+X stops the background work rather than +/// toasting, and never arms delete — so the row can settle to idle and +/// then be deleted. +#[serial_test::serial(GROK_AGENT_DASHBOARD)] +#[test] +fn dashboard_stop_bg_work_row_stops_without_arming() { + let mut app = test_app(); + let _ = dispatch_new_session_inner(&mut app, None); + let _ = dispatch_new_session_inner(&mut app, None); + let target = *app.agents.keys().next().unwrap(); + { + let agent = app.agents.get_mut(&target).unwrap(); + agent.session.session_id = Some(acp::SessionId::new("bg-loop")); + agent.session.state = AgentState::Idle; + agent.session.scheduled_tasks.insert( + "loop-1".into(), + crate::app::agent::ScheduledTaskInfo { + task_id: "loop-1".into(), + prompt: "keep going".into(), + human_schedule: "every 5m".into(), + created_at: std::time::Instant::now(), + next_fire_at: None, + tag: "loop".into(), + last_subagent_id: None, + }, + ); + } + open_dashboard(&mut app); + if let Some(d) = app.dashboard.as_mut() { + d.selected = Some(crate::views::dashboard::DashboardRowId::TopLevel(target)); + } + let effects = dispatch_dashboard_stop(&mut app); + assert!( + effects + .iter() + .any(|e| matches!(e, Effect::DeleteScheduledTask { .. })), + "Ctrl+X must stop the scheduled loop, got {effects:?}", + ); + assert!( + !effects + .iter() + .any(|e| matches!(e, Effect::DeleteSession { .. })), + "must not delete a bg-work row, got {effects:?}", + ); + let d = app.dashboard.as_ref().unwrap(); + assert!(d.delete_confirm.is_none(), "must not arm delete"); + assert!(d.error_toast.is_none(), "stopped work, so no toast"); +} +/// A row that's `Working` only because of a queued (unsent) prompt: Ctrl+X +/// drops the queue (local, no effect) rather than toasting, and never arms +/// — so the row settles to idle and can then be deleted. +#[serial_test::serial(GROK_AGENT_DASHBOARD)] +#[test] +fn dashboard_stop_queued_prompt_row_drops_queue_without_arming() { + let mut app = test_app(); + let _ = dispatch_new_session_inner(&mut app, None); + let _ = dispatch_new_session_inner(&mut app, None); + let target = *app.agents.keys().next().unwrap(); + { + let agent = app.agents.get_mut(&target).unwrap(); + agent.session.session_id = Some(acp::SessionId::new("queued")); + agent.session.state = AgentState::Idle; + agent.session.enqueue_prompt("do the thing".into()); + } + open_dashboard(&mut app); + if let Some(d) = app.dashboard.as_mut() { + d.selected = Some(crate::views::dashboard::DashboardRowId::TopLevel(target)); + } + let effects = dispatch_dashboard_stop(&mut app); + assert!( + !effects + .iter() + .any(|e| matches!(e, Effect::DeleteSession { .. })), + "must not delete a queued-prompt row, got {effects:?}", + ); + assert!( + app.agents[&target].session.pending_prompts.is_empty(), + "the queued prompt must be dropped", + ); + let d = app.dashboard.as_ref().unwrap(); + assert!(d.delete_confirm.is_none(), "must not arm delete"); + assert!(d.error_toast.is_none(), "dropped the queue, so no toast"); +} +/// The `y` / second-`[✗]` confirm re-checks deletability: a row that +/// became busy between arming and confirming must not be deleted. +#[serial_test::serial(GROK_AGENT_DASHBOARD)] +#[test] +fn dashboard_delete_confirm_rechecks_settled_row() { + let mut app = test_app(); + let _ = dispatch_new_session_inner(&mut app, None); + let _ = dispatch_new_session_inner(&mut app, None); + let target = *app.agents.keys().next().unwrap(); + { + let agent = app.agents.get_mut(&target).unwrap(); + agent.session.session_id = Some(acp::SessionId::new("recheck")); + agent.session.state = AgentState::Idle; + } + open_dashboard(&mut app); + if let Some(d) = app.dashboard.as_mut() { + d.selected = Some(crate::views::dashboard::DashboardRowId::TopLevel(target)); + } + let _ = dispatch_dashboard_stop(&mut app); + assert!(app.dashboard.as_ref().unwrap().delete_confirm.is_some()); + app.agents.get_mut(&target).unwrap().session.state = AgentState::TurnRunning; + let effects = dispatch_dashboard_delete(&mut app); + assert!( + !effects + .iter() + .any(|e| matches!(e, Effect::DeleteSession { .. })), + "a row that went busy between gestures must not delete, got {effects:?}", + ); + assert_eq!( + app.dashboard.as_ref().unwrap().error_toast.as_deref(), + Some("Stop the session before deleting"), + ); +} +/// A settled chat-conversation roster row must not arm on Ctrl+X — delete +/// isn't supported for conversations, so a confirm could never succeed. +#[serial_test::serial(GROK_AGENT_DASHBOARD)] +#[test] +fn dashboard_stop_conversation_row_does_not_arm() { + let mut app = test_app(); + let mut entry = idle_roster_entry("conv-1", "Chat row"); + entry.origin.kind = "conversation".into(); + app.dashboard_local_sessions = vec![entry]; + open_dashboard(&mut app); + if let Some(d) = app.dashboard.as_mut() { + d.selected = Some(crate::views::dashboard::DashboardRowId::Roster { + session_id: "conv-1".into(), + }); + } + assert!(dispatch_dashboard_stop(&mut app).is_empty()); + let d = app.dashboard.as_ref().unwrap(); + assert!(d.delete_confirm.is_none(), "conversation row must not arm"); + assert_eq!( + d.error_toast.as_deref(), + Some("Deleting chat conversations isn't supported yet"), + ); +} +/// A row with no session id toasts instead of emitting a delete. +#[serial_test::serial(GROK_AGENT_DASHBOARD)] +#[test] +fn dashboard_delete_top_level_without_session_id_toasts() { + let mut app = test_app(); + let _ = dispatch_new_session_inner(&mut app, None); + let _ = dispatch_new_session_inner(&mut app, None); + let target = *app.agents.keys().next().unwrap(); + app.agents.get_mut(&target).unwrap().session.session_id = None; + open_dashboard(&mut app); + if let Some(d) = app.dashboard.as_mut() { + d.selected = Some(crate::views::dashboard::DashboardRowId::TopLevel(target)); + } + let _ = dispatch_dashboard_stop(&mut app); + let effects = dispatch_dashboard_stop(&mut app); + assert!( + !effects + .iter() + .any(|e| matches!(e, Effect::DeleteSession { .. })), + "a row without a session id must not delete, got {effects:?}", + ); + assert_eq!( + app.dashboard.as_ref().unwrap().error_toast.as_deref(), + Some("No session history to delete"), + ); } /// Happy path — matching ids → no panic, queue popped. /// Also assert the response was actually sent through diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/mod.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/mod.rs index 384b625..4f640dd 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/mod.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/mod.rs @@ -29,13 +29,14 @@ use super::ctx::{find_agent_by_session_id, get_active_agent, get_active_agent_mu use super::dashboard::{ apply_pending_dispatch_config, dispatch_dashboard_attach, dispatch_dashboard_begin_rename, dispatch_dashboard_commit_rename, dispatch_dashboard_confirm_worktree, - dispatch_dashboard_create_new_agent_with_detail, dispatch_dashboard_dispatch, - dispatch_dashboard_dispatch_slash, dispatch_dashboard_overlay_cycle, - dispatch_dashboard_overlay_exit, dispatch_dashboard_overlay_stop, - dispatch_dashboard_peek_reply, dispatch_dashboard_permission_followup, - dispatch_dashboard_permission_select, dispatch_dashboard_question_answer, - dispatch_dashboard_stop, dispatch_dashboard_toggle_auto_approve, dispatch_exit_dashboard, - dispatch_open_dashboard, ensure_dashboard_state, resolve_location_input, + dispatch_dashboard_create_new_agent_with_detail, dispatch_dashboard_delete, + dispatch_dashboard_dispatch, dispatch_dashboard_dispatch_slash, + dispatch_dashboard_overlay_cycle, dispatch_dashboard_overlay_exit, + dispatch_dashboard_overlay_stop, dispatch_dashboard_peek_reply, + dispatch_dashboard_permission_followup, dispatch_dashboard_permission_select, + dispatch_dashboard_question_answer, dispatch_dashboard_stop, + dispatch_dashboard_toggle_auto_approve, dispatch_exit_dashboard, dispatch_open_dashboard, + ensure_dashboard_state, resolve_location_input, }; use super::modes::{ YOLO_ON_UNDER_PLAN_TOAST, active_agent_plan_nudge_state, dispatch_cycle_mode_and_sync, @@ -118,6 +119,16 @@ fn test_app() -> AppView { require_plan_approval: false, plan_mode: false, chat_mode: false, + #[cfg(feature = "local-workspace")] + welcome_workspace_mode: crate::views::welcome::WelcomeWorkspaceMode::Sandbox, + #[cfg(feature = "local-workspace")] + local_workspace_startup_locked: false, + #[cfg(feature = "local-workspace")] + welcome_session_local_workspace: None, + #[cfg(feature = "local-workspace")] + welcome_local_workspace_ack_pending: false, + #[cfg(feature = "local-workspace")] + welcome_history_load_as_build: false, subagents: false, ask_user: false, mouse_captured: true, @@ -206,6 +217,10 @@ fn test_app() -> AppView { welcome_privacy_banner_opt_out_rect: None, welcome_privacy_banner_terms_rect: None, welcome_privacy_banner_policy_rect: None, + #[cfg(feature = "local-workspace")] + welcome_workspace_mode_rects: Default::default(), + #[cfg(feature = "local-workspace")] + welcome_on_workspace_mode: false, welcome_toast: None, welcome_on_privacy_banner: false, welcome_on_upgrade_cta: false, @@ -228,6 +243,7 @@ fn test_app() -> AppView { session_picker_lanes: Default::default(), session_picker_detail_generation: 0, session_picker_entries_query: None, + session_picker_pending_delete: None, welcome_tick: 0, welcome_shimmer_frame: 0, startup_warnings: Vec::new(), diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/lifecycle.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/lifecycle.rs index 0d762bf..ecda77b 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/lifecycle.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/lifecycle.rs @@ -2107,6 +2107,7 @@ fn dashboard_stop_with_peek_open_moves_selection_and_peek_down_one() { let _ = dispatch_new_session_inner(&mut app, None); for (i, agent) in app.agents.values_mut().enumerate() { agent.display_name = Some(format!("agent-{i}")); + agent.session.session_id = Some(acp::SessionId::new(format!("s{i}"))); } open_dashboard(&mut app); let order = dashboard_row_order(&app); @@ -2151,17 +2152,27 @@ fn dashboard_stop_with_peek_open_moves_selection_and_peek_down_one() { other => panic!("Ctrl+X must produce DashboardStop, got {other:?}"), } } - let crate::views::dashboard::DashboardRowId::TopLevel(first_id) = first else { + let crate::views::dashboard::DashboardRowId::TopLevel(first_id) = &first else { panic!("first row should be top-level"); }; - assert!( - !app.agents.contains_key(&first_id), - "closed agent must be gone" + let session_id = app.agents[first_id] + .session + .session_id + .as_ref() + .expect("session id") + .to_string(); + let _ = dispatch_task_result( + crate::app::actions::TaskResult::DeleteSessionComplete { + source: "current".into(), + session_id, + after: crate::app::actions::AfterSessionDelete::Dashboard, + }, + &mut app, ); + assert!(!app.agents.contains_key(first_id)); assert_eq!( app.dashboard.as_ref().unwrap().selected, Some(second.clone()), - "closing with peek open should still select the next row down", ); render(&mut app); assert_eq!( @@ -2177,24 +2188,25 @@ fn dashboard_stop_with_peek_open_moves_selection_and_peek_down_one() { } /// Regression — the same Ctrl+X double-press path driven /// END-TO-END through `DashboardState::handle_input` (which the -/// existing `dashboard_stop_double_press_closes_top_level` test +/// existing `dashboard_stop_double_press_deletes_top_level` test /// bypasses by calling `dispatch_dashboard_stop` directly). /// /// Without the fix, the second `handle_input` call /// runs the top-of-`handle_key` toast/confirm clear BEFORE the /// registry resolves the key to `DashboardStop`, wiping the -/// just-armed `stop_confirm`. The dispatcher then sees a fresh -/// state and re-arms instead of closing. The session never closes +/// just-armed `delete_confirm`. The dispatcher then sees a fresh +/// state and re-arms instead of deleting. The session never deletes /// no matter how many times the user presses Ctrl+X. #[serial_test::serial(GROK_AGENT_DASHBOARD)] #[test] -fn dashboard_stop_double_press_via_handle_key_closes_top_level() { +fn dashboard_stop_double_press_via_handle_key_deletes_top_level() { use crossterm::event::{Event, KeyCode, KeyEvent, KeyModifiers}; let mut app = test_app(); let _ = dispatch_new_session_inner(&mut app, None); let _ = dispatch_new_session_inner(&mut app, None); open_dashboard(&mut app); let target = *app.agents.keys().next().unwrap(); + app.agents.get_mut(&target).unwrap().session.session_id = Some(acp::SessionId::new("s-target")); if let Some(d) = app.dashboard.as_mut() { d.selected = Some(crate::views::dashboard::DashboardRowId::TopLevel(target)); } @@ -2211,8 +2223,8 @@ fn dashboard_stop_double_press_via_handle_key_closes_top_level() { other => panic!("first Ctrl+X must produce DashboardStop, got {other:?}"), } assert!( - app.dashboard.as_ref().unwrap().stop_confirm.is_some(), - "first Ctrl+X must arm stop_confirm" + app.dashboard.as_ref().unwrap().delete_confirm.is_some(), + "first Ctrl+X must arm delete_confirm" ); let outcome2 = app .dashboard @@ -2221,14 +2233,27 @@ fn dashboard_stop_double_press_via_handle_key_closes_top_level() { .handle_input(&ctrl_x, &app.registry); match outcome2 { crate::app::app_view::InputOutcome::Action(crate::app::actions::Action::DashboardStop) => { - let _ = dispatch(crate::app::actions::Action::DashboardStop, &mut app); + let effects = dispatch(crate::app::actions::Action::DashboardStop, &mut app); + assert!(matches!(effects.last(), Some(Effect::DeleteSession { .. }))); } other => panic!("second Ctrl+X must produce DashboardStop, got {other:?}"), } - assert!( - !app.agents.contains_key(&target), - "second Ctrl+X via handle_input must close the target agent (Issue 300 regression)", + assert!(app.dashboard.as_ref().unwrap().delete_confirm.is_none()); + let session_id = app.agents[&target] + .session + .session_id + .as_ref() + .expect("session id") + .to_string(); + let _ = dispatch_task_result( + crate::app::actions::TaskResult::DeleteSessionComplete { + source: "current".into(), + session_id, + after: crate::app::actions::AfterSessionDelete::Dashboard, + }, + &mut app, ); + assert!(!app.agents.contains_key(&target)); } /// Top-level resolver round-trip via real AgentView. #[test] @@ -2274,3 +2299,884 @@ fn session_id_resolver_round_trip_subagent() { let back = resolver.to_persisted(&live).expect("must reverse"); assert_eq!(back, pid); } +#[cfg(feature = "local-workspace")] +mod welcome_workspace_mode { + use super::*; + use crate::app::session_startup::{ + LocalWorkspaceConfig, LocalWorkspaceMode, set_active_local_workspace, + }; + use crate::views::welcome::WelcomeWorkspaceMode; + #[test] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ACK)] + fn welcome_new_session_sets_own_override() { + let _ack = xai_grok_test_support::EnvGuard::set( + crate::app::session_startup::GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV, + "1", + ); + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.cwd = tmp.path().to_path_buf(); + app.welcome_workspace_mode = WelcomeWorkspaceMode::LocalWorkspace; + let _ = dispatch(Action::NewSession, &mut app); + let override_cfg = app + .welcome_session_local_workspace + .clone() + .flatten() + .expect("welcome Local must set one-shot own override"); + assert_eq!(override_cfg.mode, LocalWorkspaceMode::Own); + assert_eq!(override_cfg.cwd.as_deref(), Some(tmp.path())); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn new_session_ignores_history_bypass_for_indicator() { + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.cwd = tmp.path().to_path_buf(); + app.cwd_has_git_ancestor = false; + app.welcome_workspace_mode = WelcomeWorkspaceMode::Sandbox; + app.welcome_history_load_as_build = true; + let effects = dispatch(Action::NewSession, &mut app); + assert!( + effects + .iter() + .any(|e| matches!(e, Effect::CreateSession { .. })), + "new session must create: {effects:?}" + ); + assert!( + app.welcome_history_load_as_build, + "create must not consume history bypass (restore+load still owns it)" + ); + let agent = app.agents.values().next().expect("new agent"); + assert!(agent.chat_kind); + assert_eq!( + agent.workspace_mode, + WelcomeWorkspaceMode::Sandbox, + "create must not stamp Local from leftover history bypass" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn fork_from_welcome_with_local_selection_creates_placeholder() { + set_active_local_workspace(None).unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.welcome_workspace_mode = WelcomeWorkspaceMode::LocalWorkspace; + assert!(app.agents.is_empty()); + let effects = crate::app::dispatch::session::fork::dispatch_startup_fork_session( + &mut app, + "parent-1".into(), + None, + None, + ); + assert!( + !app.agents.is_empty(), + "fork must still create a placeholder agent" + ); + assert!( + effects + .iter() + .any(|e| matches!(e, Effect::ForkSession { .. })), + "fork effect expected: {effects:?}" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn startup_lock_prevents_sandbox_from_clearing_cli_stamp() { + let tmp = tempfile::tempdir().unwrap(); + set_active_local_workspace(Some(LocalWorkspaceConfig { + mode: LocalWorkspaceMode::Attach, + cwd: Some(tmp.path().to_path_buf()), + server_id: Some("cli-srv".into()), + })) + .unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.local_workspace_startup_locked = true; + app.welcome_workspace_mode = WelcomeWorkspaceMode::Sandbox; + app.cwd = tmp.path().to_path_buf(); + let _ = dispatch(Action::NewSession, &mut app); + let stamp = crate::app::session_startup::active_local_workspace() + .unwrap() + .expect("CLI stamp must remain"); + assert_eq!(stamp.mode, LocalWorkspaceMode::Attach); + assert!( + app.welcome_session_local_workspace.is_none(), + "locked path must not set a one-shot override" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ACK)] + fn confirm_ack_skips_reapply_and_sets_oneshot() { + let _ack = xai_grok_test_support::EnvGuard::unset( + crate::app::session_startup::GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV, + ); + let home = tempfile::tempdir().unwrap(); + let _home = + xai_grok_test_support::EnvGuard::set("GROK_HOME", home.path().to_str().unwrap()); + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.cwd = tmp.path().to_path_buf(); + app.welcome_workspace_mode = WelcomeWorkspaceMode::LocalWorkspace; + app.welcome_local_workspace_ack_pending = true; + let effects = dispatch(Action::ConfirmWelcomeLocalWorkspaceAck, &mut app); + assert!( + !app.welcome_local_workspace_ack_pending, + "confirm must clear pending" + ); + assert!( + app.welcome_session_local_workspace + .clone() + .flatten() + .is_some(), + "one-shot Own override must be set before CreateSession" + ); + assert!( + effects + .iter() + .any(|e| matches!(e, Effect::CreateSession { .. })), + "confirm must create without re-entering AwaitAck: {effects:?}" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ACK)] + fn welcome_local_worktree_always_keeps_oneshot_until_create() { + let _ack = xai_grok_test_support::EnvGuard::set( + crate::app::session_startup::GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV, + "1", + ); + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.cwd = tmp.path().to_path_buf(); + app.cwd_has_git_ancestor = true; + app.new_session_worktree_mode = crate::app::app_view::WorktreeMode::Always; + app.welcome_workspace_mode = WelcomeWorkspaceMode::LocalWorkspace; + let live = test_app_with_agent(); + let (id, agent) = live.agents.into_iter().next().unwrap(); + app.agents.insert(id, agent); + let effects = dispatch(Action::NewSession, &mut app); + assert!( + effects + .iter() + .any(|e| matches!(e, Effect::CreateWorktreeSession { .. })), + "Always worktree must emit CreateWorktreeSession: {effects:?}" + ); + assert!( + app.welcome_session_local_workspace + .clone() + .flatten() + .is_some(), + "one-shot must remain until process_effects consumes CreateWorktreeSession" + ); + assert!( + crate::app::session_startup::active_local_workspace() + .unwrap() + .is_some(), + "welcome Local stamps process-wide Own (agents map treated as stale)" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ACK)] + fn failed_worktree_create_clears_welcome_oneshot() { + let _ack = xai_grok_test_support::EnvGuard::set( + crate::app::session_startup::GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV, + "1", + ); + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.cwd = tmp.path().to_path_buf(); + app.cwd_has_git_ancestor = false; + app.welcome_workspace_mode = WelcomeWorkspaceMode::LocalWorkspace; + app.welcome_history_load_as_build = true; + let effects = dispatch( + Action::NewWorktreeSession { + load_session_id: None, + label: None, + git_ref: None, + }, + &mut app, + ); + assert!(effects.is_empty(), "expected hard-fail, got {effects:?}"); + assert!( + app.welcome_session_local_workspace.is_none(), + "failed worktree must drop one-shot so next create re-applies picker" + ); + assert!( + !app.welcome_history_load_as_build, + "failed worktree must not leak history bypass" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ACK)] + fn confirm_ack_honors_worktree_always() { + let _ack = xai_grok_test_support::EnvGuard::unset( + crate::app::session_startup::GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV, + ); + let home = tempfile::tempdir().unwrap(); + let _home = + xai_grok_test_support::EnvGuard::set("GROK_HOME", home.path().to_str().unwrap()); + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.cwd = tmp.path().to_path_buf(); + app.cwd_has_git_ancestor = true; + app.new_session_worktree_mode = crate::app::app_view::WorktreeMode::Always; + app.welcome_workspace_mode = WelcomeWorkspaceMode::LocalWorkspace; + app.welcome_local_workspace_ack_pending = true; + let effects = dispatch(Action::ConfirmWelcomeLocalWorkspaceAck, &mut app); + assert!( + effects + .iter() + .any(|e| matches!(e, Effect::CreateWorktreeSession { .. })), + "confirm must honor WorktreeMode::Always: {effects:?}" + ); + assert!( + app.welcome_session_local_workspace + .clone() + .flatten() + .is_some() + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn welcome_fetch_session_list_filters_by_workspace_mode() { + set_active_local_workspace(None).unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.welcome_workspace_mode = WelcomeWorkspaceMode::Sandbox; + let effects = dispatch(Action::FetchSessionList, &mut app); + match &effects[..] { + [Effect::FetchSessionList { kind_filter, .. }] => { + assert_eq!( + kind_filter.as_deref(), + Some(["chat".to_string()].as_slice()) + ); + } + other => panic!("expected FetchSessionList, got {other:?}"), + } + app.welcome_workspace_mode = WelcomeWorkspaceMode::LocalWorkspace; + let effects = dispatch(Action::FetchSessionList, &mut app); + match &effects[..] { + [Effect::FetchSessionList { kind_filter, .. }] => { + assert_eq!( + kind_filter.as_deref(), + Some(["build".to_string()].as_slice()) + ); + } + other => panic!("expected FetchSessionList, got {other:?}"), + } + set_active_local_workspace(None).unwrap(); + } + #[test] + fn pick_conversation_auto_switches_to_sandbox() { + set_active_local_workspace(None).unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.welcome_workspace_mode = WelcomeWorkspaceMode::LocalWorkspace; + app.session_picker_entries = Some(vec![crate::app::app_view::SessionPickerEntry { + id: "conv-1".into(), + summary: "hello".into(), + updated_at: chrono::Utc::now(), + created_at: chrono::Utc::now(), + cwd: String::new(), + hostname: None, + source: "conversation".into(), + model_id: None, + num_messages: 1, + last_active_at: None, + branch: None, + repo_name: String::new(), + worktree_label: None, + card_detail: None, + }]); + let effects = dispatch(Action::PickSession(0), &mut app); + assert_eq!(app.welcome_workspace_mode, WelcomeWorkspaceMode::Sandbox); + assert!( + app.welcome_session_local_workspace.is_none(), + "conversation pick must drop (not force-clear) local one-shot" + ); + assert!( + effects.iter().any(|e| matches!( + e, + Effect::LoadSession { + chat_kind: true, + .. + } + )), + "conversation must load as chat: {effects:?}" + ); + assert!(!app.welcome_history_load_as_build); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn pick_local_disk_auto_switches_to_local_and_bypasses_chat_refusal() { + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.cwd = tmp.path().to_path_buf(); + app.welcome_workspace_mode = WelcomeWorkspaceMode::Sandbox; + let sess_dir = super::super::super::plant_local_build_session(tmp.path(), "build-1"); + app.session_picker_entries = Some(vec![crate::app::app_view::SessionPickerEntry { + id: "build-1".into(), + summary: "local work".into(), + updated_at: chrono::Utc::now(), + created_at: chrono::Utc::now(), + cwd: tmp.path().display().to_string(), + hostname: None, + source: "local".into(), + model_id: None, + num_messages: 1, + last_active_at: None, + branch: None, + repo_name: String::new(), + worktree_label: None, + card_detail: None, + }]); + let effects = dispatch(Action::PickSession(0), &mut app); + assert_eq!( + app.welcome_workspace_mode, + WelcomeWorkspaceMode::LocalWorkspace + ); + assert!(app.chat_mode, "sticky --chat remains"); + assert!( + effects.iter().any(|e| matches!( + e, + Effect::LoadSession { + chat_kind: false, + .. + } + )), + "local-disk pick must load as build: {effects:?}" + ); + assert!( + app.welcome_history_load_as_build, + "bypass stays until process_effects LoadSession" + ); + let agent = app.agents.values().next().expect("placeholder agent"); + assert!( + agent.chat_kind, + "sticky --chat keeps agent.chat_kind for already-open focus matching" + ); + assert_eq!( + agent.workspace_mode, + WelcomeWorkspaceMode::LocalWorkspace, + "Local UX is the workspace_mode indicator" + ); + let _ = std::fs::remove_dir_all(sess_dir); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn pick_local_disk_in_worktree_sets_history_bypass() { + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.cwd = tmp.path().to_path_buf(); + app.cwd_has_git_ancestor = true; + app.welcome_workspace_mode = WelcomeWorkspaceMode::Sandbox; + app.session_picker_entries = Some(vec![crate::app::app_view::SessionPickerEntry { + id: "build-wt".into(), + summary: "local work".into(), + updated_at: chrono::Utc::now(), + created_at: chrono::Utc::now(), + cwd: tmp.path().display().to_string(), + hostname: None, + source: "local".into(), + model_id: None, + num_messages: 1, + last_active_at: None, + branch: None, + repo_name: String::new(), + worktree_label: None, + card_detail: None, + }]); + let _ = dispatch(Action::PickSessionInWorktree(0), &mut app); + assert_eq!( + app.welcome_workspace_mode, + WelcomeWorkspaceMode::LocalWorkspace + ); + assert!( + app.welcome_history_load_as_build, + "worktree pick of build row must set history bypass" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn pick_in_worktree_resume_skips_local_ack() { + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.cwd = tmp.path().to_path_buf(); + app.cwd_has_git_ancestor = true; + app.welcome_workspace_mode = WelcomeWorkspaceMode::LocalWorkspace; + app.session_picker_entries = Some(vec![crate::app::app_view::SessionPickerEntry { + id: "build-wt".into(), + summary: "local work".into(), + updated_at: chrono::Utc::now(), + created_at: chrono::Utc::now(), + cwd: tmp.path().display().to_string(), + hostname: None, + source: "local".into(), + model_id: None, + num_messages: 1, + last_active_at: None, + branch: None, + repo_name: String::new(), + worktree_label: None, + card_detail: None, + }]); + let effects = dispatch(Action::PickSessionInWorktree(0), &mut app); + assert!( + !app.welcome_local_workspace_ack_pending, + "worktree resume must not block on Local ACK" + ); + assert!( + effects + .iter() + .any(|e| matches!(e, Effect::CreateWorktreeSession { .. })), + "worktree resume must create worktree without ACK: {effects:?}" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ACK)] + fn pick_in_worktree_no_git_clears_history_bypass() { + let _ack = xai_grok_test_support::EnvGuard::set( + crate::app::session_startup::GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV, + "1", + ); + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.cwd = tmp.path().to_path_buf(); + app.cwd_has_git_ancestor = false; + app.welcome_workspace_mode = WelcomeWorkspaceMode::Sandbox; + app.session_picker_entries = Some(vec![crate::app::app_view::SessionPickerEntry { + id: "build-wt".into(), + summary: "local work".into(), + updated_at: chrono::Utc::now(), + created_at: chrono::Utc::now(), + cwd: tmp.path().display().to_string(), + hostname: None, + source: "local".into(), + model_id: None, + num_messages: 1, + last_active_at: None, + branch: None, + repo_name: String::new(), + worktree_label: None, + card_detail: None, + }]); + let effects = dispatch(Action::PickSessionInWorktree(0), &mut app); + assert!( + effects.is_empty(), + "no-git worktree must hard-fail: {effects:?}" + ); + assert!( + !app.welcome_history_load_as_build, + "no-git worktree fail must not leak history bypass" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn cli_lock_still_sets_history_bypass_without_rewriting_mode() { + let tmp = tempfile::tempdir().unwrap(); + set_active_local_workspace(Some(LocalWorkspaceConfig { + mode: LocalWorkspaceMode::Attach, + cwd: Some(tmp.path().to_path_buf()), + server_id: Some("cli-srv".into()), + })) + .unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.cwd = tmp.path().to_path_buf(); + app.local_workspace_startup_locked = true; + app.welcome_workspace_mode = WelcomeWorkspaceMode::Sandbox; + let sess_dir = super::super::super::plant_local_build_session(tmp.path(), "build-lock"); + app.session_picker_entries = Some(vec![crate::app::app_view::SessionPickerEntry { + id: "build-lock".into(), + summary: "local work".into(), + updated_at: chrono::Utc::now(), + created_at: chrono::Utc::now(), + cwd: tmp.path().display().to_string(), + hostname: None, + source: "local".into(), + model_id: None, + num_messages: 1, + last_active_at: None, + branch: None, + repo_name: String::new(), + worktree_label: None, + card_detail: None, + }]); + let effects = dispatch(Action::PickSession(0), &mut app); + assert_eq!( + app.welcome_workspace_mode, + WelcomeWorkspaceMode::Sandbox, + "CLI lock must not rewrite welcome mode from Sandbox" + ); + assert!( + app.welcome_history_load_as_build, + "CLI lock must still set local-disk load bypass" + ); + assert!( + effects.iter().any(|e| matches!( + e, + Effect::LoadSession { + chat_kind: false, + .. + } + )), + "locked local-disk pick must still load: {effects:?}" + ); + let _ = std::fs::remove_dir_all(sess_dir); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn failed_local_pick_clears_history_bypass() { + set_active_local_workspace(None).unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.welcome_workspace_mode = WelcomeWorkspaceMode::Sandbox; + app.session_picker_entries = Some(vec![crate::app::app_view::SessionPickerEntry { + id: "missing-build".into(), + summary: "gone".into(), + updated_at: chrono::Utc::now(), + created_at: chrono::Utc::now(), + cwd: String::new(), + hostname: None, + source: "local".into(), + model_id: None, + num_messages: 1, + last_active_at: None, + branch: None, + repo_name: String::new(), + worktree_label: None, + card_detail: None, + }]); + let effects = dispatch(Action::PickSession(0), &mut app); + assert!(effects.is_empty()); + assert!( + !app.welcome_history_load_as_build, + "failed/no-op pick must not leak bypass" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn deferred_history_bypass_survives_startup_gate() { + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.trust_state = crate::app::app_view::TrustState::Pending { + workspace: tmp.path().to_path_buf(), + }; + app.welcome_history_load_as_build = true; + let effects = dispatch(Action::LoadSession("sid".into(), None, false), &mut app); + assert!(effects.is_empty()); + assert!( + !app.welcome_history_load_as_build, + "live flag moved onto deferred startup" + ); + assert!(app.deferred_startup.history_load_as_build); + let effects = finish_trust(&mut app); + assert!( + app.welcome_history_load_as_build, + "drain must re-apply bypass before LoadSession" + ); + assert!( + effects.iter().any(|e| matches!( + e, + Effect::LoadSession { + chat_kind: false, + .. + } + )), + "deferred drain must emit LoadSession: {effects:?}" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn cli_lock_conversation_pick_does_not_rewrite_sandbox_mode() { + let tmp = tempfile::tempdir().unwrap(); + set_active_local_workspace(Some(LocalWorkspaceConfig { + mode: LocalWorkspaceMode::Attach, + cwd: Some(tmp.path().to_path_buf()), + server_id: Some("cli-srv".into()), + })) + .unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.local_workspace_startup_locked = true; + app.welcome_workspace_mode = WelcomeWorkspaceMode::Sandbox; + app.session_picker_entries = Some(vec![crate::app::app_view::SessionPickerEntry { + id: "conv-lock".into(), + summary: "hello".into(), + updated_at: chrono::Utc::now(), + created_at: chrono::Utc::now(), + cwd: String::new(), + hostname: None, + source: "conversation".into(), + model_id: None, + num_messages: 1, + last_active_at: None, + branch: None, + repo_name: String::new(), + worktree_label: None, + card_detail: None, + }]); + let effects = dispatch(Action::PickSession(0), &mut app); + assert_eq!( + app.welcome_workspace_mode, + WelcomeWorkspaceMode::Sandbox, + "CLI lock must not auto-switch welcome mode on conversation pick" + ); + assert!(!app.welcome_history_load_as_build); + assert!( + effects.iter().any(|e| matches!( + e, + Effect::LoadSession { + chat_kind: true, + .. + } + )), + "conversation must still load: {effects:?}" + ); + let agent = app.agents.values().next().expect("agent"); + assert_eq!( + agent.workspace_mode, + WelcomeWorkspaceMode::Sandbox, + "conversation without session-local intent → Sandbox (not Local·CLI)" + ); + assert!( + !agent.workspace_mode_cli_locked, + "CLI lock must not badge conversation LoadSession" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn session_restore_failed_clears_history_bypass() { + set_active_local_workspace(None).unwrap(); + let mut app = test_app_with_agent(); + app.chat_mode = true; + app.welcome_history_load_as_build = true; + let id = AgentId(0); + let effects = dispatch( + Action::TaskComplete(TaskResult::SessionRestoreFailed { + agent_id: id, + error: "boom".into(), + }), + &mut app, + ); + assert!(effects.is_empty()); + assert!( + !app.welcome_history_load_as_build, + "failed restore must not leak bypass into the next load" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn restore_and_load_sets_local_workspace_indicator() { + set_active_local_workspace(None).unwrap(); + let mut app = test_app(); + app.chat_mode = true; + app.active_view = ActiveView::Welcome; + app.welcome_workspace_mode = WelcomeWorkspaceMode::Sandbox; + app.session_picker_entries = Some(vec![crate::app::app_view::SessionPickerEntry { + id: "remote-1".into(), + summary: "remote row".into(), + updated_at: chrono::Utc::now(), + created_at: chrono::Utc::now(), + cwd: "/other".into(), + hostname: None, + source: "remote".into(), + model_id: None, + num_messages: 1, + last_active_at: None, + branch: None, + repo_name: String::new(), + worktree_label: None, + card_detail: None, + }]); + let effects = dispatch(Action::PickSession(0), &mut app); + assert!( + effects + .iter() + .any(|e| matches!(e, Effect::RestoreAndLoadSession { .. })), + "remote pick must restore: {effects:?}" + ); + assert!( + app.welcome_history_load_as_build, + "bypass kept until follow-up LoadSession" + ); + let agent = app.agents.values().next().expect("restore placeholder"); + assert_eq!( + agent.workspace_mode, + WelcomeWorkspaceMode::LocalWorkspace, + "restore placeholder must show Local indicator" + ); + set_active_local_workspace(None).unwrap(); + } + #[test] + fn history_build_bypass_only_applies_to_load_session_batch() { + use crate::app::event_loop::welcome_history_build_bypass_applies; + assert!(!welcome_history_build_bypass_applies(&[], true)); + assert!(!welcome_history_build_bypass_applies( + &[Effect::FetchSessionList { + query: None, + seq: 0, + kind_filter: None, + }], + true + )); + assert!(welcome_history_build_bypass_applies( + &[Effect::LoadSession { + agent_id: AgentId(0), + session_id: "s".into(), + session_cwd: None, + chat_kind: false, + }], + true + )); + assert!(welcome_history_build_bypass_applies( + &[Effect::RestoreAndLoadSession { + agent_id: AgentId(0), + session_id: "s".into(), + session_cwd: "/tmp".into(), + }], + true + )); + assert!(welcome_history_build_bypass_applies( + &[Effect::CreateWorktreeSession { + agent_id: AgentId(0), + load_session_id: Some("s".into()), + label: None, + git_ref: None, + model_id: None, + preferred_session_id: None, + chat_kind: false, + }], + true + )); + assert!( + crate::app::event_loop::welcome_history_build_bypass_consume( + &[Effect::CreateWorktreeSession { + agent_id: AgentId(0), + load_session_id: Some("s".into()), + label: None, + git_ref: None, + model_id: None, + preferred_session_id: None, + chat_kind: false, + }], + true + ), + "worktree-resume batch consumes bypass (single-batch create)" + ); + assert!( + !crate::app::event_loop::welcome_history_build_bypass_consume( + &[Effect::RestoreAndLoadSession { + agent_id: AgentId(0), + session_id: "s".into(), + session_cwd: "/tmp".into(), + }], + true + ), + "restore-only batch keeps bypass for follow-up LoadSession" + ); + assert!( + crate::app::event_loop::welcome_history_build_bypass_consume( + &[Effect::LoadSession { + agent_id: AgentId(0), + session_id: "s".into(), + session_cwd: None, + chat_kind: false, + }], + true + ) + ); + } + #[test] + fn fetch_session_list_kind_filter_only_on_welcome_chat() { + set_active_local_workspace(None).unwrap(); + let mut welcome = test_app(); + welcome.chat_mode = true; + welcome.active_view = ActiveView::Welcome; + welcome.welcome_workspace_mode = WelcomeWorkspaceMode::LocalWorkspace; + match &dispatch(Action::FetchSessionList, &mut welcome)[..] { + [Effect::FetchSessionList { kind_filter, .. }] => { + assert_eq!( + kind_filter.as_deref(), + Some(["build".to_string()].as_slice()) + ); + } + other => panic!("{other:?}"), + } + let mut in_session = test_app_with_agent(); + in_session.chat_mode = true; + match &dispatch(Action::FetchSessionList, &mut in_session)[..] { + [Effect::FetchSessionList { kind_filter, .. }] => { + assert!(kind_filter.is_none()) + } + other => panic!("{other:?}"), + } + set_active_local_workspace(None).unwrap(); + } + #[test] + fn in_session_new_does_not_clear_process_stamp() { + let tmp = tempfile::tempdir().unwrap(); + set_active_local_workspace(Some(LocalWorkspaceConfig { + mode: LocalWorkspaceMode::Own, + cwd: Some(tmp.path().to_path_buf()), + server_id: None, + })) + .unwrap(); + let mut app = test_app_with_agent(); + app.chat_mode = true; + app.cwd = tmp.path().to_path_buf(); + app.welcome_workspace_mode = WelcomeWorkspaceMode::Sandbox; + let _ = dispatch(Action::NewSession, &mut app); + assert!( + crate::app::session_startup::active_local_workspace() + .unwrap() + .is_some(), + "in-session /new must not clear the process stamp" + ); + set_active_local_workspace(None).unwrap(); + } +} diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/load.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/load.rs index 1a2131d..0d33da6 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/load.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/session/load.rs @@ -1647,7 +1647,7 @@ fn chat_mode_debounce_expiry_fetches_current_and_drops_stale() { assert!( matches!( &effects[..], - [Effect::FetchSessionList { query: Some(q), seq: 1 }] if q == "abc" + [Effect::FetchSessionList { query: Some(q), seq: 1, .. }] if q == "abc" ), "current debounce expiry must fetch with the query, got {effects:?}" ); @@ -1879,7 +1879,7 @@ fn chat_mode_force_search_fetches_immediately_and_empty_query_unfilters() { assert!( matches!( &effects[..], - [Effect::FetchSessionList { query: Some(q), seq: 1 }] if q == "abc" + [Effect::FetchSessionList { query: Some(q), seq: 1, .. }] if q == "abc" ), "forced search must fetch without debouncing, got {effects:?}" ); @@ -1894,7 +1894,8 @@ fn chat_mode_force_search_fetches_immediately_and_empty_query_unfilters() { &effects[..], [Effect::FetchSessionList { query: None, - seq: 2 + seq: 2, + .. }] ), "cleared query must refetch the unfiltered list immediately (no debounce), got {effects:?}" @@ -2575,7 +2576,8 @@ fn build_mode_rapid_plain_fetches_keep_last_write_wins() { &effects[..], [Effect::FetchSessionList { query: None, - seq: 0 + seq: 0, + .. }] ), "Build-mode plain fetch must not bump the seq, got {effects:?}" @@ -2634,7 +2636,8 @@ fn plain_picker_fetch_carries_no_query_and_bumps_seq() { &effects[..], [Effect::FetchSessionList { query: None, - seq: 2 + seq: 2, + .. }] ), "picker fetch must be unfiltered and supersede the search, got {effects:?}" diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/task_result.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/task_result.rs index 941371a..267f231 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/tests/task_result.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/tests/task_result.rs @@ -1694,7 +1694,11 @@ fn delete_session_complete_removes_only_matching_source_and_id() { .active_modal .as_mut() { - *pending_delete = Some(("local".into(), "s1".into(), "/r".into())); + *pending_delete = Some(crate::views::session_picker::PendingDelete { + source: "local".into(), + session_id: "s1".into(), + cwd: "/r".into(), + }); } let _ = dispatch_task_result( diff --git a/crates/codegen/xai-grok-pager/src/app/dispatch/turn.rs b/crates/codegen/xai-grok-pager/src/app/dispatch/turn.rs index 0e9d135..4c1ae20 100644 --- a/crates/codegen/xai-grok-pager/src/app/dispatch/turn.rs +++ b/crates/codegen/xai-grok-pager/src/app/dispatch/turn.rs @@ -93,10 +93,13 @@ pub(super) fn dispatch_cancel_turn(app: &mut AppView) -> Vec { rewind_if_pristine: false, }]; } - if !agent.session.state.is_turn_running() { + if !agent.session.state.is_turn_running() && !agent.session.state.is_compact_running() { return vec![]; } - if let Some(stop) = resolved_pref { + if agent.session.state.is_compact_running() { + // No subagent picker for `/compact` — just stop the generation. + resolved_pref.or(Some(true)) + } else if let Some(stop) = resolved_pref { Some(stop) } else { // Check all running subagents, not just those from the current turn. @@ -183,6 +186,22 @@ pub(super) fn do_cancel_turn(app: &mut AppView, cancel_subagents: bool) -> Vec Vec { }] } -// TODO: Add dispatch_cancel_command() once xai-grok-shell supports proper -// server-side cancellation for /compact. Currently, the compaction handler -// uses spawn_local with no cancellation token, and blindly replaces the -// conversation history when done — so prompts sent after a client-side -// cancel would be lost. - // TaskResult handlers. pub(super) fn handle_bg_task_killed( diff --git a/crates/codegen/xai-grok-pager/src/app/effects/helpers.rs b/crates/codegen/xai-grok-pager/src/app/effects/helpers.rs index f0dc48f..326047b 100644 --- a/crates/codegen/xai-grok-pager/src/app/effects/helpers.rs +++ b/crates/codegen/xai-grok-pager/src/app/effects/helpers.rs @@ -271,6 +271,9 @@ pub(crate) struct SessionFlags { /// Mutual exclusivity with Build plan profiles: profiles are omitted and a /// warn is logged when plan flags are also set (K12). pub chat_mode: bool, + /// Local-workspace stamp for ACP `_meta` (scrub still strips envId / Direct hub). + #[cfg(feature = "local-workspace")] + pub local_workspace: Option, /// Effective screen mode label (`ScreenMode::meta_label`), stamped into /// every `PromptRequest._meta.screenMode` for minimal-vs-regular usage /// telemetry. `None` (key omitted) only under `Default` in tests; real @@ -327,6 +330,10 @@ impl SessionFlags { } if self.chat_mode { meta.insert("x.ai/session".into(), serde_json::json!({ "kind": "chat" })); + #[cfg(feature = "local-workspace")] + if let Some(ref lw) = self.local_workspace { + stamp_local_workspace_meta(&mut meta, lw); + } } if !self.ask_user { meta.insert("askUserQuestion".into(), serde_json::json!(false)); @@ -346,33 +353,107 @@ impl SessionFlags { /// /// `x.ai/cloud_existing_workspace` is intentionally omitted: scrub keeps it /// iff `x.ai/local_workspace.mode == "attach"`. +#[allow(dead_code)] pub(super) const CHAT_FORBIDDEN_WORKSPACE_BIND_KEYS: &[&str] = &[ "envId", "x.ai/cloud_server_id", ]; +/// FS-only tool ids for local existing workspace (chat attach/own). +#[cfg(feature = "local-workspace")] +pub(super) const LOCAL_WORKSPACE_FS_ONLY_TOOL_IDS: &[&str] = &[ + "workspace.fs_list", + "workspace.fs_exists", + "workspace.fs_read_file", + "workspace.fs_write_file", + "workspace.fs_delete_file", + "workspace.put_files", + "workspace.get_files", +]; /// Stamp `_meta["x.ai/session"].kind = "chat"` and strip Build `agentProfile` (K12). pub(super) fn apply_chat_kind_meta(meta: &mut Option) { let obj = meta.get_or_insert_with(acp::Meta::new); obj.insert("x.ai/session".into(), serde_json::json!({ "kind": "chat" })); obj.remove("agentProfile"); } +/// Stamp chat+local intent. Attach also stamps `x.ai/cloud_existing_workspace`. +/// Own leaves `server_id` unset — shell supervisor mints before handshake. +/// +/// Never stamps `envId` or `x.ai/cloud_server_id`. +#[cfg(feature = "local-workspace")] +pub(super) fn stamp_local_workspace_meta( + meta: &mut serde_json::Map, + cfg: &crate::app::session_startup::LocalWorkspaceConfig, +) { + use crate::app::session_startup::LocalWorkspaceMode; + let mut local = serde_json::Map::new(); + let mode = match cfg.mode { + LocalWorkspaceMode::Attach => "attach", + LocalWorkspaceMode::Own => "own", + }; + local.insert("mode".into(), serde_json::json!(mode)); + if let Some(ref sid) = cfg.server_id { + local.insert("server_id".into(), serde_json::json!(sid)); + } + if let Some(ref cwd) = cfg.cwd { + local + .insert("cwd".into(), serde_json::json!(cwd.to_string_lossy().into_owned())); + } + meta.insert("x.ai/local_workspace".into(), serde_json::Value::Object(local)); + tracing::info!( + target: crate::views::welcome::workspace_mode::WORKSPACE_MODE_LOG, + event = "acp_meta_stamped", + mode, + server_id = cfg.server_id.as_deref(), + cwd = cfg.cwd.as_ref().map(|p| p.display().to_string()), + "stamped x.ai/local_workspace onto session meta" + ); + if cfg.mode == LocalWorkspaceMode::Attach && let Some(ref sid) = cfg.server_id { + let mut existing = serde_json::Map::new(); + existing.insert("server_id".into(), serde_json::json!(sid)); + if let Some(ref cwd) = cfg.cwd { + existing + .insert( + "cwd".into(), + serde_json::json!(cwd.to_string_lossy().into_owned()), + ); + } + meta.insert( + "x.ai/cloud_existing_workspace".into(), + serde_json::Value::Object(existing), + ); + } +} +/// Apply [`stamp_local_workspace_meta`] onto optional ACP meta. +#[cfg(feature = "local-workspace")] +pub(super) fn apply_local_workspace_meta( + meta: &mut Option, + cfg: &crate::app::session_startup::LocalWorkspaceConfig, +) { + let obj = meta.get_or_insert_with(acp::Meta::new); + stamp_local_workspace_meta(obj, cfg); +} /// Shared chat create/load/worktree meta finalize: kind + local stamp + scrub. pub(super) fn finalize_chat_session_meta( meta: &mut Option, is_chat_path: bool, - #[allow(unused_variables)] + #[cfg_attr(not(feature = "local-workspace"), allow(unused_variables))] session_flags: &SessionFlags, ) { if !is_chat_path { return; } apply_chat_kind_meta(meta); + #[cfg(feature = "local-workspace")] + if let Some(ref lw) = session_flags.local_workspace { + apply_local_workspace_meta(meta, lw); + } scrub_chat_workspace_bind_meta(meta); } /// Remove client workspace-bind keys from chat create/load meta (defense in depth). /// /// Narrow scrub exception: keep `x.ai/cloud_existing_workspace` when local -/// intent is attach. Never keep `envId` or Direct hub `x.ai/cloud_server_id`. +/// intent is **attach**. Own stamps intent only (shell mints `server_id`). +/// Never keep `envId` or Direct hub `x.ai/cloud_server_id`. pub(super) fn scrub_chat_workspace_bind_meta(meta: &mut Option) { let Some(obj) = meta.as_mut() else { return; @@ -380,10 +461,80 @@ pub(super) fn scrub_chat_workspace_bind_meta(meta: &mut Option) { for key in CHAT_FORBIDDEN_WORKSPACE_BIND_KEYS { obj.remove(*key); } + #[cfg(feature = "local-workspace")] + { + let allow_existing_attach = obj + .get("x.ai/local_workspace") + .and_then(|v| v.get("mode")) + .and_then(|m| m.as_str()) == Some("attach"); + if !allow_existing_attach { + obj.remove("x.ai/cloud_existing_workspace"); + } + } { obj.remove("x.ai/cloud_existing_workspace"); } } +/// Params for shell ACP `x.ai/session/add_local_workspace`. +/// +/// v1 surface is **shell ACP-only** (no pager slash/command wiring). Pager +/// dogfood / headless clients call the extension directly with this payload. +/// No remove path until session end. +#[cfg(feature = "local-workspace")] +#[allow(dead_code)] +pub(crate) fn mid_session_add_local_workspace_params( + session_id: &str, + cfg: &crate::app::session_startup::LocalWorkspaceConfig, +) -> serde_json::Value { + let mut meta = serde_json::Map::new(); + stamp_local_workspace_meta(&mut meta, cfg); + let mut opt = Some(meta); + scrub_chat_workspace_bind_meta(&mut opt); + serde_json::json!({ + "sessionId": session_id, + "meta": opt.unwrap_or_default(), + }) +} +/// Fail closed on operator attestation outside the FS-only allowlist. +/// `None` / empty attested set → uncheckable → refuse. Live server is not probed. +#[cfg(feature = "local-workspace")] +pub(crate) fn reject_non_fs_only_advertised_tools( + advertised_tool_ids: Option<&[&str]>, +) -> Result<(), String> { + let Some(ids) = advertised_tool_ids else { + return Err( + "operator attestation GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS is unset \ + (uncheckable); refuse attach. Live workspace_server was not inspected — set \ + the env to a comma-separated FS-only catalog." + .into(), + ); + }; + if ids.is_empty() { + return Err( + "operator attestation GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS is empty \ + (uncheckable); refuse attach. Live workspace_server was not inspected." + .into(), + ); + } + let forbidden: Vec<&str> = ids + .iter() + .copied() + .filter(|id| !LOCAL_WORKSPACE_FS_ONLY_TOOL_IDS.contains(id)) + .collect(); + if forbidden.is_empty() { + Ok(()) + } else { + Err( + format!( + "operator attestation lists tools outside the FS-only allowlist: {}. \ + Live workspace_server was not inspected. Fix \ + GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS or restart workspace_server \ + with --require-explicit-toolset and an FS-only catalog.", + forbidden.join(", ") + ), + ) + } +} /// Metadata returned from effect execution so the event loop can patch /// state that requires a spawned task handle (e.g., auth AbortHandle). #[derive(Default)] diff --git a/crates/codegen/xai-grok-pager/src/app/effects/mod.rs b/crates/codegen/xai-grok-pager/src/app/effects/mod.rs index c17efca..8875841 100644 --- a/crates/codegen/xai-grok-pager/src/app/effects/mod.rs +++ b/crates/codegen/xai-grok-pager/src/app/effects/mod.rs @@ -18,6 +18,8 @@ pub(crate) use helpers::{ EffectMeta, RestoreProgressMsg, SessionFlags, persist_permission_mode_and_notify, persist_setting, sanitize_user_error, }; +#[cfg(feature = "local-workspace")] +pub(crate) use helpers::reject_non_fs_only_advertised_tools; use helpers::*; use std::path::{Path, PathBuf}; use agent_client_protocol as acp; @@ -707,7 +709,7 @@ pub(crate) fn execute( } }); } - Effect::FetchSessionList { query, seq } => { + Effect::FetchSessionList { query, seq, kind_filter } => { let tx = acp_tx.clone(); let cwd = cwd.to_path_buf(); tasks @@ -721,6 +723,19 @@ pub(crate) fn execute( } else { params["allowRelax"] = serde_json::Value::Bool(true); } + if let Some(kinds) = &kind_filter { + params["_meta"] = serde_json::json!({ + "x.ai/facetFilters": { "kind": kinds }, + }); + tracing::info!( + target: "grok.pager.workspace_mode", + event = "session_list_fetch", + kind_filter = ?kinds, + query = ?query, + seq, + "FetchSessionList with kind facet filter" + ); + } let request = acp::ExtRequest::new( "x.ai/session/list", serde_json::value::to_raw_value(¶ms) @@ -3543,6 +3558,7 @@ pub(crate) fn execute( } Effect::SendBtw { agent_id, session_id, question, minimal_request_id } => { let tx = acp_tx.clone(); + let is_api_key_auth = session_flags.is_api_key_auth; tasks .spawn(async move { let request = acp::ExtRequest::new( @@ -3577,9 +3593,7 @@ pub(crate) fn execute( Err(e) => { TaskResult::BtwResponse { agent_id, - result: Err( - sanitize_user_error(&format!("side question failed: {e}")), - ), + result: Err(format_acp_error(&e, is_api_key_auth)), minimal_request_id, } } diff --git a/crates/codegen/xai-grok-pager/src/app/effects/tests.rs b/crates/codegen/xai-grok-pager/src/app/effects/tests.rs index 25d8bd5..04472a7 100644 --- a/crates/codegen/xai-grok-pager/src/app/effects/tests.rs +++ b/crates/codegen/xai-grok-pager/src/app/effects/tests.rs @@ -1523,6 +1523,7 @@ async fn fetch_session_list_pushes_query_and_echoes_seq() { let mut tasks = run(Effect::FetchSessionList { query: Some("hit".into()), seq: 7, + kind_filter: None, }); match tasks.join_next().await.expect("task").expect("no panic") { TaskResult::SessionListLoaded { sessions, scope, seq, query, .. } => { @@ -1539,6 +1540,7 @@ async fn fetch_session_list_pushes_query_and_echoes_seq() { let mut tasks = run(Effect::FetchSessionList { query: None, seq: 8, + kind_filter: None, }); match tasks.join_next().await.expect("task").expect("no panic") { TaskResult::SessionListLoaded { scope, seq, query, .. } => { @@ -1554,6 +1556,7 @@ async fn fetch_session_list_pushes_query_and_echoes_seq() { let mut tasks = run(Effect::FetchSessionList { query: Some("fail-me".into()), seq: 9, + kind_filter: None, }); match tasks.join_next().await.expect("task").expect("no panic") { TaskResult::SessionListFailed { error, seq, query } => { @@ -1589,6 +1592,50 @@ async fn fetch_session_list_pushes_query_and_echoes_seq() { assert_eq!(captured[2]["query"], "fail-me"); } #[tokio::test] +async fn fetch_session_list_sends_kind_facet_filter() { + use std::sync::{Arc, Mutex}; + use xai_acp_lib::AcpAgentMessage; + let captured: Arc>> = Arc::default(); + let captured_for_task = captured.clone(); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + tokio::spawn(async move { + while let Some(msg) = rx.recv().await { + if let AcpAgentMessage::ExtMethod(args) = msg { + let params: serde_json::Value = serde_json::from_str( + args.request.params.get(), + ) + .expect("params JSON"); + captured_for_task.lock().unwrap().push(params); + let body = serde_json::json!({ "result": { "sessions": [] } }); + let raw = serde_json::value::RawValue::from_string(body.to_string()) + .expect("ser"); + let _ = args.response_tx.send(Ok(acp::ExtResponse::new(Arc::from(raw)))); + } + } + }); + let (progress_tx, _progress_rx) = tokio::sync::mpsc::unbounded_channel(); + let mut tasks = JoinSet::new(); + execute( + Effect::FetchSessionList { + query: None, + seq: 1, + kind_filter: Some(vec!["build".into()]), + }, + &mut tasks, + &tx, + Path::new("."), + &SessionFlags::default(), + &progress_tx, + ); + let _ = tasks.join_next().await; + let captured = captured.lock().unwrap(); + assert_eq!(captured.len(), 1); + assert_eq!( + captured[0]["_meta"]["x.ai/facetFilters"]["kind"], + serde_json::json!(["build"]) + ); +} +#[tokio::test] async fn fetch_workflows_list_sends_session_id() { use std::sync::{Arc, Mutex}; use xai_acp_lib::AcpAgentMessage; @@ -1979,7 +2026,7 @@ fn to_meta_chat_mode_stamps_kind_and_omits_agent_profile() { ..Default::default() }; let meta = flags.to_meta().expect("chat_mode must emit meta"); - assert_eq!(meta["x.ai/session"] ["kind"], "chat"); + assert_eq!(meta["x.ai/session"]["kind"], "chat"); assert!( meta.get("agentProfile").is_none(), "K12: chat mode must omit Build agentProfile" @@ -2006,7 +2053,7 @@ fn load_meta_chat_kind_alone_stamps_kind_and_strips_profile() { scrub_chat_workspace_bind_meta(&mut meta); } let meta = meta.expect("chat_kind must produce meta"); - assert_eq!(meta["x.ai/session"] ["kind"], "chat"); + assert_eq!(meta["x.ai/session"]["kind"], "chat"); assert!( meta.get("agentProfile").is_none(), "entry chat_kind must strip Build agentProfile" @@ -2040,7 +2087,7 @@ fn chat_create_meta_never_includes_workspace_bind_keys_when_cloud_fields_set() { apply_chat_kind_meta(&mut meta); scrub_chat_workspace_bind_meta(&mut meta); let meta = meta.expect("chat create must emit meta"); - assert_eq!(meta["x.ai/session"] ["kind"], "chat"); + assert_eq!(meta["x.ai/session"]["kind"], "chat"); assert_chat_meta_has_no_workspace_bind_keys( &serde_json::Value::Object(meta.clone()), ); @@ -2064,11 +2111,165 @@ fn chat_load_meta_never_includes_workspace_bind_keys() { } scrub_chat_workspace_bind_meta(&mut meta); let meta = meta.expect("chat load must emit meta"); - assert_eq!(meta["x.ai/session"] ["kind"], "chat"); + assert_eq!(meta["x.ai/session"]["kind"], "chat"); assert_chat_meta_has_no_workspace_bind_keys( &serde_json::Value::Object(meta.clone()), ); } +/// Attach stamp keeps existing workspace + local intent; envId / Direct hub stay stripped. +#[cfg(feature = "local-workspace")] +#[test] +fn scrub_chat_workspace_matrix_attach_exception() { + use crate::app::session_startup::{LocalWorkspaceConfig, LocalWorkspaceMode}; + let mut meta = Some(acp::Meta::new()); + { + let obj = meta.as_mut().unwrap(); + obj.insert("envId".into(), serde_json::json!("env-x")); + obj.insert("x.ai/cloud_server_id".into(), serde_json::json!("hub-x")); + obj.insert( + "x.ai/cloud_existing_workspace".into(), + serde_json::json!({"server_id": "srv-x", "cwd": "/ws"}), + ); + } + scrub_chat_workspace_bind_meta(&mut meta); + let scrubbed = meta.as_ref().unwrap(); + assert!(scrubbed.get("envId").is_none()); + assert!(scrubbed.get("x.ai/cloud_server_id").is_none()); + assert!(scrubbed.get("x.ai/cloud_existing_workspace").is_none()); + let mut meta = Some(acp::Meta::new()); + apply_local_workspace_meta( + &mut meta, + &LocalWorkspaceConfig { + mode: LocalWorkspaceMode::Attach, + cwd: Some(std::path::PathBuf::from("/tmp/repo")), + server_id: Some("srv-dogfood".into()), + }, + ); + { + let obj = meta.as_mut().unwrap(); + obj.insert("envId".into(), serde_json::json!("env-must-go")); + obj.insert("x.ai/cloud_server_id".into(), serde_json::json!("hub-must-go")); + } + scrub_chat_workspace_bind_meta(&mut meta); + let scrubbed = meta.as_ref().unwrap(); + assert!(scrubbed.get("envId").is_none(), "envId must stay scrubbed"); + assert!( + scrubbed.get("x.ai/cloud_server_id").is_none(), + "Direct hub must stay scrubbed" + ); + assert_eq!( + scrubbed["x.ai/cloud_existing_workspace"]["server_id"], + "srv-dogfood" + ); + assert_eq!(scrubbed["x.ai/local_workspace"]["mode"], "attach"); + assert_eq!(scrubbed["x.ai/local_workspace"]["server_id"], "srv-dogfood"); + assert_eq!(scrubbed["x.ai/local_workspace"]["cwd"], "/tmp/repo"); +} +#[cfg(feature = "local-workspace")] +#[test] +fn to_meta_chat_attach_stamps_local_and_existing() { + use crate::app::session_startup::{LocalWorkspaceConfig, LocalWorkspaceMode}; + let flags = SessionFlags { + chat_mode: true, + local_workspace: Some(LocalWorkspaceConfig { + mode: LocalWorkspaceMode::Attach, + cwd: Some(std::path::PathBuf::from("/tmp/repo")), + server_id: Some("srv-1".into()), + }), + ..Default::default() + }; + let meta = flags.to_meta().expect("meta"); + assert_eq!(meta["x.ai/session"]["kind"], "chat"); + assert_eq!(meta["x.ai/local_workspace"]["mode"], "attach"); + assert_eq!(meta["x.ai/cloud_existing_workspace"]["server_id"], "srv-1"); + assert!(meta.get("envId").is_none()); + assert!(meta.get("x.ai/cloud_server_id").is_none()); +} +#[cfg(feature = "local-workspace")] +#[test] +fn to_meta_chat_own_stamps_intent_without_existing() { + use crate::app::session_startup::{LocalWorkspaceConfig, LocalWorkspaceMode}; + let flags = SessionFlags { + chat_mode: true, + local_workspace: Some(LocalWorkspaceConfig { + mode: LocalWorkspaceMode::Own, + cwd: Some(std::path::PathBuf::from("/tmp/repo-own")), + server_id: None, + }), + ..Default::default() + }; + let meta = flags.to_meta().expect("meta"); + assert_eq!(meta["x.ai/local_workspace"]["mode"], "own"); + assert_eq!(meta["x.ai/local_workspace"]["cwd"], "/tmp/repo-own"); + assert!(meta["x.ai/local_workspace"].get("server_id").is_none()); + assert!( + meta.get("x.ai/cloud_existing_workspace").is_none(), + "own must not stamp existing; shell mints server_id" + ); + assert!(meta.get("envId").is_none()); +} +#[cfg(feature = "local-workspace")] +#[test] +fn mid_session_add_params_scrub_envid() { + use crate::app::session_startup::{LocalWorkspaceConfig, LocalWorkspaceMode}; + let params = mid_session_add_local_workspace_params( + "sess-1", + &LocalWorkspaceConfig { + mode: LocalWorkspaceMode::Attach, + cwd: Some(std::path::PathBuf::from("/tmp/repo")), + server_id: Some("srv-add".into()), + }, + ); + assert_eq!(params["sessionId"], "sess-1"); + assert_eq!(params["meta"]["x.ai/local_workspace"]["mode"], "attach"); + assert_eq!( + params["meta"]["x.ai/cloud_existing_workspace"]["server_id"], + "srv-add" + ); + assert!(params["meta"].get("envId").is_none()); +} +#[cfg(feature = "local-workspace")] +#[test] +fn reject_non_fs_only_advertised_tools_matrix() { + let fs_only = ["workspace.fs_list", "workspace.fs_read_file", "workspace.put_files"]; + assert!(reject_non_fs_only_advertised_tools(Some(&fs_only[..])).is_ok()); + assert!( + reject_non_fs_only_advertised_tools(None) + .unwrap_err() + .contains("uncheckable") + ); + assert!( + reject_non_fs_only_advertised_tools(Some(&[][..])) + .unwrap_err() + .contains("empty") + ); + let with_exec = ["workspace.fs_list", "workspace.bash", "terminal.exec"]; + let err = reject_non_fs_only_advertised_tools(Some(&with_exec[..])).unwrap_err(); + assert!(err.contains("FS-only"), "{err}"); + assert!(err.contains("workspace.bash"), "{err}"); + assert!(err.contains("terminal.exec"), "{err}"); +} +#[cfg(feature = "local-workspace")] +#[test] +fn finalize_chat_session_meta_stamps_attach_on_worktree_path() { + use crate::app::session_startup::{LocalWorkspaceConfig, LocalWorkspaceMode}; + let flags = SessionFlags { + chat_mode: false, + local_workspace: Some(LocalWorkspaceConfig { + mode: LocalWorkspaceMode::Attach, + cwd: Some(std::path::PathBuf::from("/tmp/repo")), + server_id: Some("srv-wt".into()), + }), + ..Default::default() + }; + let mut meta = flags.to_meta(); + finalize_chat_session_meta(&mut meta, true, &flags); + let meta = meta.expect("meta"); + assert_eq!(meta["x.ai/session"]["kind"], "chat"); + assert_eq!(meta["x.ai/local_workspace"]["mode"], "attach"); + assert_eq!(meta["x.ai/cloud_existing_workspace"]["server_id"], "srv-wt"); + assert!(meta.get("envId").is_none()); +} #[test] fn to_meta_yolo_suppresses_auto_mode() { let flags = SessionFlags { diff --git a/crates/codegen/xai-grok-pager/src/app/event_loop.rs b/crates/codegen/xai-grok-pager/src/app/event_loop.rs index 98cbd5b..562c880 100644 --- a/crates/codegen/xai-grok-pager/src/app/event_loop.rs +++ b/crates/codegen/xai-grok-pager/src/app/event_loop.rs @@ -4,27 +4,34 @@ //! management is delegated to [`AppView`]. The event loop only handles //! IO plumbing: terminal events, ACP channel, spawned task results, //! animation ticks, and hot-reloadable config changes. + +use std::time::Duration; + +use anyhow::Context as _; +use crossterm::event::{Event, KeyCode, KeyEvent, KeyEventKind, KeyModifiers}; +use tokio::task::JoinSet; +use tokio::time::{Instant, sleep_until}; + +use crate::appearance::ConfigWatcher; +use crate::client_identity::{PAGER_CLIENT_TYPE, PAGER_CLIENT_VERSION}; +use crate::theme::system_appearance::{self, SystemAppearanceWatcher}; +use crate::theme::{Theme, ThemeKind, cache as theme_cache}; + +use agent_client_protocol as acp; +use xai_acp_lib::acp_send; + use super::actions::{Action, Effect, TaskResult}; use super::app_view::{ ActiveView, AppView, AuthState, InputOutcome, PasteProvenance, TrustState, VoiceState, }; use super::{PagerArgs, PagerTerminal, acp_handler, dispatch, effects}; -use crate::appearance::ConfigWatcher; -use crate::client_identity::{PAGER_CLIENT_TYPE, PAGER_CLIENT_VERSION}; -use crate::theme::system_appearance::{self, SystemAppearanceWatcher}; -use crate::theme::{Theme, ThemeKind, cache as theme_cache}; -use agent_client_protocol as acp; -use anyhow::Context as _; -use crossterm::event::{Event, KeyCode, KeyEvent, KeyEventKind, KeyModifiers}; -use std::time::Duration; -use tokio::task::JoinSet; -use tokio::time::{Instant, sleep_until}; -use xai_acp_lib::acp_send; + #[derive(Clone, Debug, PartialEq)] pub(super) struct TimedInputEvent { pub(super) event: Event, pub(super) arrived_at: std::time::Instant, } + impl TimedInputEvent { fn now(event: Event) -> Self { Self { @@ -33,6 +40,7 @@ impl TimedInputEvent { } } } + /// Values resolved before `init_terminal` and consumed by the event loop. /// /// All fields must be computed while stdin is still in cooked mode and @@ -48,6 +56,7 @@ pub(crate) struct TerminalState { /// its OSC 11 fallback reads stdin and competes with the input reader. pub initial_theme: ThemeKind, } + /// Result of the event loop run. pub(crate) struct RunResult { pub exit_info: Option, @@ -56,6 +65,7 @@ pub(crate) struct RunResult { /// terminal restore. See `/minimal` and `/fullscreen`. pub relaunch: Option, } + /// In-flight reconnect re-initialization, tied to the agents whose reload /// windows it opened so completion lands on them even if the user switches /// views (or closes one) while the re-init runs. @@ -67,6 +77,7 @@ struct ReconnectReinit { /// Reconnect generation that opened the reload windows. generation: u64, } + /// Result of a reconnect re-initialization task. struct ReinitOutcome { /// Whether initialize/authenticate succeeded; when false no load was @@ -74,6 +85,7 @@ struct ReinitOutcome { init_ok: bool, loads: Vec, } + /// Per-agent `session/load` outcome from the re-init task. struct AgentLoadOutcome { agent_id: super::agent::AgentId, @@ -88,6 +100,7 @@ struct AgentLoadOutcome { /// describes a runtime the new actor will not use. scheduler_background_loops: Option, } + /// Fields of the reconnect `session/load`, derived from the agent being /// reloaded. `None` when the agent has no session yet. struct ReconnectLoadPlan { @@ -101,6 +114,7 @@ struct ReconnectLoadPlan { /// and full-replays when it doesn't. meta: serde_json::Value, } + fn restore_dashboard_peek_before_reload( dashboard: &mut Option, agents: &mut indexmap::IndexMap, @@ -109,6 +123,7 @@ fn restore_dashboard_peek_before_reload( dashboard.restore_peek_viewport(agents); } } + fn plan_reconnect_load( agent: &super::agent_view::AgentView, fallback_cwd: &std::path::Path, @@ -120,6 +135,12 @@ fn plan_reconnect_load( agent.session.cwd.clone() }; let yolo = agent.session.is_yolo(); + // Set BOTH yoloMode and autoMode explicitly. The leader's capability injection + // only fills ABSENT keys, so omitting autoMode here lets a stale launch-time + // `ClientCapabilities.auto_mode` re-enable Auto after the user left it (e.g. + // Shift+Tab to Ask). Auto is per-agent (symmetric with yolo) — derive it from + // this agent's own `auto_mode` so a background tab reconnects with ITS mode, + // not the active tab's global `current_ui` mirror. let auto = super::dispatch::effective_auto(yolo, agent.session.is_auto()); let mut meta = serde_json::json!({ "yoloMode": yolo, "autoMode": auto }); if let Some(ref cursor) = agent.last_seen_event_id { @@ -131,6 +152,7 @@ fn plan_reconnect_load( meta, }) } + /// Resolve the two post-reconnect restore outcomes from the per-agent /// `session/load` results. /// @@ -158,6 +180,7 @@ fn reconnect_restore_outcome( && active_agent_id.is_some_and(|aid| pending_agent_ids.contains(&aid) && load_ok(&aid)); (all_restored, active_restored) } + /// Compute the folder-trust verdict for the session cwd and seed /// [`AppView::trust_state`]. Pager-side mirror of the agent's resolve: read the /// local store, scan for repo-local code-exec config, and run the pure @@ -176,19 +199,30 @@ fn seed_trust_state( TrustOutcome, decide, decide_inputs_with_interactive, feature_enabled, }; use xai_grok_workspace::trust::workspace_key; + let feature = feature_enabled(remote); if !feature { app.trust_state = TrustState::Done; return; } + + // The cwd the user launched in == the process cwd == `app.cwd` (set at + // construction), matching the `--trust` grant's `std::env::current_dir()`. let cwd = app.cwd.clone(); let key = workspace_key(&cwd); + // Reuse the canonical gather (store trust + repo-config scan) but pass the + // pager's stdin-only interactivity: the TUI prompts via the rendered + // question + crossterm keyboard, NOT stderr (the pager redirects native + // stderr at startup, so the engine's `stdin && stderr` would be false here + // and the question would never show). TTY stdin => user can answer; + // otherwise fail closed (no prompt). let inputs = decide_inputs_with_interactive(&cwd, &key, std::io::stdin().is_terminal()); app.trust_state = match decide(feature, &inputs) { TrustOutcome::Prompt => TrustState::Pending { workspace: key }, TrustOutcome::Trusted | TrustOutcome::Untrusted => TrustState::Done, }; } + /// Pause terminal input and wait up to `timeout` for the reader to acknowledge. /// Returns with the pause still asserted; the handoff owner resumes the reader. fn park_input_reader( @@ -197,6 +231,8 @@ fn park_input_reader( timeout: Duration, ) -> bool { use std::sync::atomic::Ordering; + // Storing `reader_parked = false` before `input_paused = true` is + // intentionally ordered to prevent accepting a stale parked acknowledgement. reader_parked.store(false, Ordering::Release); input_paused.store(true, Ordering::Release); let deadline = std::time::Instant::now() + timeout; @@ -205,6 +241,7 @@ fn park_input_reader( } reader_parked.load(Ordering::Acquire) } + /// Suspend the TUI, let a blocking child own the tty, then restore it. /// /// Input is parked before the asynchronous frame writer is drained with a @@ -242,6 +279,9 @@ fn suspend_for_child( return Err(error); } } + + // Pre-child cursor probe (minimal only — minimal's startup already proved + // this terminal answers CPR). Reader is parked, so the reply is ours. let pre_cursor = screen_mode .is_minimal() .then(|| crossterm::cursor::position().ok()) @@ -259,17 +299,24 @@ fn suspend_for_child( let _ = crossterm::execute!(stderr, crossterm::terminal::EnterAlternateScreen); }); } + // Discard child-exit ANSI query replies (DA/DSR/cursor reports) the terminal + // buffered; reader is parked, so the main thread is the only crossterm caller. while crossterm::event::poll(Duration::from_millis(0)).unwrap_or(false) { let _ = crossterm::event::read(); } + // Post-child cursor probe: `Some` iff the child left the cursor somewhere + // other than where it found it; restore_after_child uses that to re-anchor + // minimal mode after main-screen output. let moved_cursor = pre_cursor.and_then(|pre| { let post = crossterm::cursor::position().ok()?; (post != pre).then_some(post) }); + // Only the pre-park race can reach this channel; later input stays in the tty. while input_rx.try_recv().is_ok() {} input_paused.store(false, Ordering::Release); Ok(moved_cursor) } + /// Coalesces draw requests, gates in-flight frames, and owns draw cadence. #[derive(Debug)] struct Presenter { @@ -279,6 +326,7 @@ struct Presenter { last_draw_at: Instant, draw_scheduled_at: Option, } + impl Presenter { fn new() -> Self { Self { @@ -289,6 +337,7 @@ impl Presenter { draw_scheduled_at: None, } } + fn acknowledge(&mut self, sequence: u64) { if self .in_flight_target @@ -297,6 +346,7 @@ impl Presenter { self.in_flight_target = None; } } + fn try_present( &mut self, queued_before: u64, @@ -315,10 +365,12 @@ impl Presenter { } true } + fn request(&mut self, force_full_repaint: bool) { self.dirty = true; self.force_full_repaint |= force_full_repaint; } + /// Request now when cadence permits; otherwise schedule the earliest draw. fn request_throttled(&mut self, now: Instant, min_draw_interval: Duration) -> bool { if now.duration_since(self.last_draw_at) < min_draw_interval { @@ -330,10 +382,12 @@ impl Presenter { self.request(false); true } + fn mark_drawn(&mut self, now: Instant) { self.last_draw_at = now; self.draw_scheduled_at = None; } + fn present_if_dirty(&mut self, app: &mut AppView, terminal: &mut PagerTerminal) { let sync = terminal.backend_mut().writer_mut().writer_sync().clone(); let queued_before = sync.queued(); @@ -351,6 +405,7 @@ impl Presenter { self.mark_drawn(Instant::now()); } } + fn request_presentation( &mut self, app: &mut AppView, @@ -361,21 +416,26 @@ impl Presenter { self.present_if_dirty(app, terminal); } } + fn writer_event_sequence(event: crate::render::draw::WriterEvent) -> std::io::Result { match event { crate::render::draw::WriterEvent::Written(sequence) => Ok(sequence), crate::render::draw::WriterEvent::Failed(error) => Err(error), } } + const SUSPEND_RETRY_DELAY: Duration = Duration::from_millis(250); + fn suspend_retry_ready(retry_after: Option, now: Instant) -> bool { retry_after.is_none_or(|deadline| now >= deadline) } + #[derive(Debug, Default)] struct SuspendWaitReports { editor_reported: bool, pager_reported: bool, } + impl SuspendWaitReports { fn reset_missing(&mut self, editor_pending: bool, pager_pending: bool) { if !editor_pending { @@ -386,6 +446,7 @@ impl SuspendWaitReports { } } } + /// Arm the deferred retry and return whether this pending handoff needs feedback. fn defer_suspend_retry( retry_after: &mut Option, @@ -398,13 +459,16 @@ fn defer_suspend_retry( *wait_reported = true; should_report } + const EDITOR_SUSPEND_WAIT: &str = "Editor is waiting for a safe terminal handoff"; const TRANSCRIPT_SUSPEND_WAIT: &str = "Transcript is waiting for a safe terminal handoff"; + #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum SuspendWaitSink { Toast, SystemBlock, } + fn suspend_wait_sink(screen_mode: crate::app::ScreenMode) -> SuspendWaitSink { if screen_mode.is_minimal() { SuspendWaitSink::SystemBlock @@ -412,6 +476,7 @@ fn suspend_wait_sink(screen_mode: crate::app::ScreenMode) -> SuspendWaitSink { SuspendWaitSink::Toast } } + /// Report a handoff wait through the sink visible in the current screen mode. /// The caller deduplicates reports across retries per handoff request. fn report_suspend_wait(app: &mut AppView, message: &str) { @@ -433,9 +498,12 @@ fn report_suspend_wait(app: &mut AppView, message: &str) { } } } + fn requeue_after_suspend_timeout(pending: &mut Option, request: T) { + // The child never started, so preserve the one-shot request. *pending = Some(request); } + /// Restore presentation after a child releases the tty. /// /// A cat-style child leaves minimal mode's cursor below appended main-screen @@ -454,6 +522,7 @@ fn restore_after_child( let screen = terminal.last_known_area(); let cur = terminal.viewport_area(); let vh = cur.height.max(1).min(screen.height.max(1)); + // Buffered append stays ordered before the gated repaint. let _ = terminal.backend_mut().append_lines(vh.saturating_sub(1)); let available = screen.height.saturating_sub(y).saturating_sub(1); let top = y.saturating_sub(vh.saturating_sub(1).saturating_sub(available)); @@ -464,6 +533,7 @@ fn restore_after_child( }); } } + /// Consume a pending `$EDITOR` / `$PAGER` suspend request, if any. /// /// Called at the top of every event-loop iteration because any select arm can @@ -488,11 +558,17 @@ fn run_pending_suspends( if !suspend_retry_ready(*suspend_retry_after, Instant::now()) { return Ok(()); } + // The gate is consumed before any blocking park/drain attempt. A timeout + // must arm a fresh deadline before this function returns. if !editor_pending && !pager_pending { *suspend_retry_after = None; return Ok(()); } *suspend_retry_after = None; + + // $EDITOR suspend: leave alt screen, disable raw mode, spawn + // editor, wait for exit, then restore. Preparation materializes prompt + // drafts only immediately before this safe terminal handoff. if let Some(request) = app.pending_editor.take() { let retry_request = request.clone(); match crate::app::external_editor::prepare(app, request) { @@ -533,6 +609,9 @@ fn run_pending_suspends( Err(error) => return Err(error.into()), }; crate::app::external_editor::finish(app, prepared, editor_result); + // The child owned the screen; re-anchor if it printed inline, and + // repaint the full viewport rather than diffing against a screen + // state we can no longer vouch for. restore_after_child(terminal, app.screen_mode, moved_cursor); presenter.request_presentation(app, terminal, true); suspend_wait_reports.editor_reported = false; @@ -548,6 +627,10 @@ fn run_pending_suspends( } } } + + // /transcript suspend: open the rendered transcript in $PAGER, + // then restore and delete the temp file. Shares the editor's + // suspend/restore dance (reader park, raw mode, alt screen). if let Some(path) = app.pending_pager_path.take() { let ansi = std::mem::take(&mut app.pending_pager_ansi); let pager = std::env::var("PAGER") @@ -561,9 +644,15 @@ fn run_pending_suspends( reader_parked, input_rx, || { + // $PAGER may carry flags (e.g. "less -R"); split on + // whitespace so program + args are both honored. let mut parts = pager.split_whitespace(); if let Some(prog) = parts.next() { let mut args: Vec = parts.map(str::to_string).collect(); + // An ANSI transcript (minimal full view) needs + // `less` to interpret raw control codes, else the + // colors show as literal escapes. Add `-R` when + // using less and it isn't already requested. let is_less = std::path::Path::new(prog) .file_name() .and_then(|n| n.to_str()) @@ -579,6 +668,10 @@ fn run_pending_suspends( { args.push("-R".to_string()); } + // Open the transcript at its END: minimal's prompt sits at + // the bottom of the conversation, so the pager starts where + // the user already is (`g` jumps back to the top). less-only + // like `-R` — other $PAGERs may not understand `+G`. if ansi && is_less && !args.iter().any(|a| a == "+G") { args.push("+G".to_string()); } @@ -607,12 +700,16 @@ fn run_pending_suspends( Err(error) => return Err(error.into()), }; let _ = std::fs::remove_file(&path); + // The pager owned the screen; re-anchor if it printed inline (cat) and + // repaint the full viewport rather than diffing against a screen state + // we can no longer vouch for. restore_after_child(terminal, app.screen_mode, moved_cursor); presenter.request_presentation(app, terminal, true); suspend_wait_reports.pager_reported = false; } Ok(()) } + /// Run the main event loop until quit. /// /// Returns a [`RunResult`] with optional exit info (for the resume hint) @@ -636,10 +733,16 @@ pub(crate) async fn run( >, mut writer_event_rx: tokio::sync::mpsc::UnboundedReceiver, ) -> anyhow::Result { + // Initialize tracing capture. The channel `rx` will be wired to a + // TracingModel (and ultimately a tracing pane) once integrated. + // For now we drain-and-discard in `AppView::tick()` to avoid unbounded + // memory growth. if args.log_sampling { + // SAFETY: called before any threads are spawned by init_tracing. unsafe { std::env::set_var("GROK_LOG_SAMPLING", "1") }; } let tracing_handle = crate::tracing::init_tracing(); + crate::unified_log::init(connection.tx.clone()); crate::unified_log::info("pager started", None, None); let mut app = AppView::new( @@ -648,10 +751,22 @@ pub(crate) async fn run( connection.available_commands, ); app.tracing_rx = Some(tracing_handle.rx); + // Startup terminal height for the auto-compact derivation; kept fresh by + // `Event::Resize` from here on. 0 (probe failure) never forces compact. app.last_known_terminal_rows = crossterm::terminal::size().map(|(_, r)| r).unwrap_or(0); + // Leader mode: a live `leader_status_rx` means the pager is connected via a + // leader. The dashboard itself is NOT gated on this flag (it renders local + // sessions regardless); `leader_mode` only controls whether we additionally + // poll the leader roster (see the roster-poll arm below). app.leader_mode = connection.leader_status_rx.is_some(); app.screen_mode = term_state.screen_mode; + // `AppView::new` precedes the terminal's resolved screen mode. Rebuild the + // registry at this I/O boundary; the later config-aware rebuild preserves + // this mode while adding the optional mouse-reporting action. app.registry = crate::actions::ActionRegistry::defaults_for(term_state.screen_mode); + // Agent/dashboard prompts pick the mode up at their creation sites + // (`apply_app_scoped_gates` / `ensure_dashboard_state`); the welcome prompt + // already exists, so inject here. app.welcome_prompt.set_screen_mode(term_state.screen_mode); if app.screen_mode.is_minimal() && term_state.relaunched_into_minimal { app.minimal_state.welcome_pending = true; @@ -670,6 +785,8 @@ pub(crate) async fn run( remote_permission_mode, ); app.default_yolo = launch_yolo.yolo; + // Gated launch-auto (CLI `--permission-mode auto` or config). Hoisted so it can + // be re-applied after `load_initial_ui_config()` replaces `current_ui` below. let launch_auto = xai_grok_shell::util::config::effective_auto_for_launch( args.yolo, args.permission_mode_flag.as_deref(), @@ -678,19 +795,27 @@ pub(crate) async fn run( if launch_auto { app.current_ui.permission_mode = Some("auto".into()); } + // One effective-config read for launch-mode ownership + the display + // resolve below (the launch resolvers above keep their own internal read). let launch_effective_ui = xai_grok_shell::config::load_effective_config() .ok() .and_then(|root| root.get("ui").cloned()); + // Soft-default owns the mode only when neither CLI nor effective TOML + // claimed it; while owned, `settings/update` pushes may re-arm it. let cli_owns_mode = args.yolo || args.permission_mode_flag.is_some(); let toml_owns_mode = launch_effective_ui .as_ref() .and_then(xai_grok_shell::util::config::permission_mode_from_ui_if_set) .is_some(); app.permission_mode_from_soft_default = !cli_owns_mode && !toml_owns_mode; + // Cached pin snapshot gating dispatch's runtime always-approve toggles. A + // mid-session pin change is missed here, but only cosmetically: the agent's + // permission manager re-clamps yolo authoritatively at decision time. app.yolo_policy_block = launch_yolo.policy_block; if let Some(warning) = launch_yolo.blocked_warning { tracing::warn!("{warning}"); crate::unified_log::warn(warning, None, None); + // Consumed by `switch_to_agent` once the first agent view opens. app.yolo_launch_block_notice = Some(warning); } app.require_plan_approval = xai_grok_shell::util::config::load_require_plan_approval(); @@ -698,6 +823,18 @@ pub(crate) async fn run( app.subagents = !args.no_subagents; app.ask_user = !args.no_ask_user; app.chat_mode = args.chat(); + #[cfg(feature = "local-workspace")] + { + let stamp = crate::app::session_startup::active_local_workspace() + .ok() + .flatten(); + app.local_workspace_startup_locked = stamp.is_some(); + if app.local_workspace_startup_locked { + app.welcome_workspace_mode = + crate::views::welcome::workspace_mode::mode_from_active_stamp(stamp.as_ref()); + crate::views::welcome::workspace_mode::log_cli_lock_applied(app.welcome_workspace_mode); + } + } app.restore_code = args.restore_code.then_some(true); if let Some(ref agent) = args.agent { match crate::headless::resolve_agent_arg(agent) { @@ -758,6 +895,7 @@ pub(crate) async fn run( .as_ref() .and_then(|s| s.privacy_banner_reshow_days) }); + // Local dismiss timestamp for the coding-data privacy banner. app.privacy_banner_acked = xai_grok_shell::config::load_from_disk() .ok() .and_then(|root| { @@ -768,6 +906,7 @@ pub(crate) async fn run( app.plugin_cta_enabled = xai_grok_config::env_bool("GROK_PLUGIN_CTA") .or_else(|| remote_settings.as_ref().and_then(|s| s.plugin_cta)) .unwrap_or(false); + // Voice is applied after auth_meta so API-key detection is accurate. app.session_picker_grouped = std::env::var("GROK_SESSION_PICKER_GROUPED") .ok() .and_then(|v| match v.as_str() { @@ -788,12 +927,19 @@ pub(crate) async fn run( .unwrap_or(true); app.cancel_rewind_enabled = connection.cancel_rewind_enabled; apply_session_recap_available(&mut app, connection.session_recap_available); + + // Preserve auth methods so logout→re-login works without restarting. app.auth_methods = connection.auth_methods.clone(); + + // Seed auth state from ACP connection metadata. + // --force-login overrides: show the login screen even when credentials exist. let force_login = args.force_login && !connection.auth_methods.is_empty(); let needs_interactive_login = connection.needs_login || force_login; if needs_interactive_login { app.welcome_prompt_focused = false; + if connection.needs_login { + // Normal path: use the metadata from startup_auth_metadata() app.login_label = connection.login_label; app.login_method_id = connection.login_method_id; app.auth_start_mode = match connection.auth_start_mode { @@ -801,6 +947,7 @@ pub(crate) async fn run( crate::acp::AuthStartMode::Command => super::app_view::AuthMode::Command, }; } else { + // --force-login: find the grok.com method from the advertised list let grok_com = connection .auth_methods .iter() @@ -820,20 +967,31 @@ pub(crate) async fn run( super::app_view::AuthMode::Pending }; } else { + // No grok.com method available, use the first method as fallback let first = &connection.auth_methods[0]; app.login_label = Some(first.name().to_string()); app.login_method_id = Some(first.id().clone()); app.auth_start_mode = super::app_view::AuthMode::Pending; } } + + // Skip the login splash screen — auto-trigger login immediately + // by reusing dispatch_login. Effects are stashed and drained after + // the initial render so the user sees the auth UI right away. + // Empty auth_methods (preferred_method pin with no credentials) is + // fail-closed: do not invent grok.com / auto-start OIDC. tracing::info!( method_id = ?app.login_method_id, methods_empty = connection.auth_methods.is_empty(), "auto-triggering login at startup" ); } + // 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 { if connection.auth_methods.is_empty() { + // preferred_method pin unavailable — no advertised method to start. app.auth_state = super::app_view::AuthState::Pending { error: Some( xai_grok_shell::agent::auth_method::PREFERRED_API_KEY_UNAVAILABLE.to_string(), @@ -846,22 +1004,29 @@ pub(crate) async fn run( } else { vec![] }; + app.has_external_auth_provider = crate::slash::commands::usage::detect_external_auth_provider(&app.auth_methods); + if let Some(meta) = connection.auth_meta.as_ref() { match serde_json::from_value::(meta.clone()) { Ok(auth_meta) => app.apply_auth_meta(&auth_meta), Err(e) => tracing::warn!("failed to deserialize auth_meta: {e}"), } } else { + // No cached session — check if the API key is the active credential. app.is_api_key_auth = app.auth_methods.iter().any(|m| { m.id().0.as_ref() == xai_grok_shell::agent::auth_method::XAI_API_KEY_METHOD_ID }); + // No AuthMeta on this path — API keys / external auth have no + // consumer billing surface. External auth also hides `/usage`. if app.is_api_key_auth || app.has_external_auth_provider { app.usage_visible = false; app.sync_billing_surface_to_agents(); } } + + // After auth so API-key + managed policy resolve correctly. let voice_mode_enabled = crate::app::resolve_voice_mode_live( remote_settings.as_ref().and_then(|s| s.voice_mode_enabled), app.is_api_key_auth, @@ -871,18 +1036,31 @@ pub(crate) async fn run( app.voice_ui_active = false; } app.apply_voice_mode_enabled(voice_mode_enabled); + + // 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 persisted per-ID hidden state app.hidden_announcement_ids = xai_grok_announcements::read_hidden_announcement_ids().await; + + // Load config layers once, resolve announcements, tips, and feature flags. let requirements = xai_grok_shell::config::load_merged_requirements(); let user_config = xai_grok_shell::config::load_from_disk().ok(); let managed_config = xai_grok_shell::config::load_managed_config().ok(); + + // Full merge when every layer parses; partial merge below if any layer fails. let effective_config = match xai_grok_shell::config::load_effective_config() { Ok(raw) => Some(raw), Err(e) => { @@ -900,11 +1078,15 @@ pub(crate) async fn run( codex: compat.codex.sessions, cursor: compat.cursor.sessions, }; + + // Load notification config from [ui.notifications] in config.toml. if let Some(ref raw) = effective_config { app.notification_service = crate::notifications::NotificationService::new( crate::notifications::load_notification_config(raw), ); if let Some(table) = raw.as_table() { + // Voice inherits the same resolved endpoints base as chat + // (config > GROK_XAI_API_BASE_URL env > default). let endpoints_base = xai_grok_shell::agent::config::EndpointsConfig::from_config_value(raw) .xai_api_base_url; @@ -912,17 +1094,26 @@ pub(crate) async fn run( xai_grok_voice::VoiceConfig::from_config_table(table, Some(&endpoints_base)); } } + // Stamp request-identity headers so the STT handshake attributes voice usage + // to grok-cli server-side (mirrors sampler / imagine). Done after + // `from_config_table` — which yields a fresh config with these + // `#[serde(skip)]` fields defaulted to empty — and unconditionally, so they + // apply even when there is no `[voice]` table (or no config at all). app.voice_config.client_identifier = crate::client_identity::HEADLESS_CLIENT_TYPE.to_string(); app.voice_config.user_agent = crate::client_identity::client_user_agent(); + app.zdr_access_enabled = xai_grok_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( xai_grok_shell::util::config::resolve_show_thinking_blocks( requirements.as_ref(), @@ -950,22 +1141,32 @@ pub(crate) async fn run( ) .value, ); + + // Pre-arrival seed only. The authoritative per-session value rides the + // `session/new` / `session/load` response, but `/loop` can be reached from + // the session-less dashboard and from a session whose response has not + // landed yet; both need an answer now, and this is the same resolver the + // shell runs at spawn, so the seed agrees with the flag as it stands today. app.scheduler_background_loops_seed = xai_grok_shell::util::config::resolve_scheduler_background_loops( remote_settings .as_ref() .and_then(|s| s.scheduler_background_loops), ); + 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 xai_grok_shell::util::config::{ resolve_announcements, resolve_slash_command_tags, resolve_tips, }; + let remote_announcements = remote_settings .as_ref() .and_then(|s| s.announcements.as_deref()); @@ -982,6 +1183,7 @@ pub(crate) async fn run( app.announcement = app.active_announcements.get(idx).cloned(); } app.sync_session_announcement_slash_gate(); + let remote_tips = remote_settings.as_ref().and_then(|s| s.tips.as_deref()); app.tips = resolve_tips( requirements.as_ref(), @@ -989,10 +1191,14 @@ pub(crate) async fn run( managed_config.as_ref(), remote_tips, ); + if !app.tips.is_empty() { let grok_home = xai_grok_tools::util::grok_home::grok_home(); app.tip = xai_grok_shell::util::tips::pick_and_advance(&app.tips, &grok_home); } + + // Slash-command dropdown tags: remote base, local [slash_command_tags] + // wins per key. Mutate the shared map in place so every adopter sees it. let remote_slash_tags = remote_settings .as_ref() .and_then(|s| s.slash_command_tags.as_ref()); @@ -1000,6 +1206,7 @@ pub(crate) async fn run( let tags_config = effective_config.as_ref().unwrap_or(&empty_toml); *app.command_tags.borrow_mut() = resolve_slash_command_tags(tags_config, remote_slash_tags); } + let hints = xai_grok_shell::util::config::resolve_hints( effective_config.as_ref(), requirements.as_ref(), @@ -1007,12 +1214,21 @@ pub(crate) async fn run( managed_config.as_ref(), ); app.project_picker_disabled = hints.project_picker_disabled; + // Per-tip contextual hints resolve from `[ui.contextual_hints]` (loaded into + // `app.current_ui` further below) + the remote tier; the resolve + prompt + // propagation happen after `current_ui` is hydrated. app.remote_contextual_hints = remote_settings .as_ref() .and_then(|s| s.contextual_hints.clone()); app.new_session_worktree_mode = hints.new_session_worktree_mode.into(); app.fork_worktree_mode = hints.fork_worktree_mode.into(); + // Ephemeral-tip seen counts are intentionally NOT hydrated: the cap is + // per-session (in-memory `app.tip_seen_counts`), so each run starts fresh. + + // Cache whether cwd is inside a git repo (avoids repeated stat() in draw). app.cwd_has_git_ancestor = app.cwd.ancestors().any(|p| p.join(".git").exists()); + + // Probe / auto-cadence / terminal telemetry — see `display_refresh_startup`. let motion = super::display_refresh_startup::start( requirements.as_ref(), user_config.as_ref(), @@ -1021,6 +1237,10 @@ pub(crate) async fn run( ); let min_draw_interval = motion.min_draw_interval; let scroll_cadence = motion.scroll_cadence; + + // Collect structured startup warnings from the terminal diagnostics engine. + // These are stored on AppView and rendered as a dismissible in-app banner + // when the user enters an agent session. { let ctx = crate::terminal::terminal_context(); let query = crate::diagnostics::probes::LiveTmuxProbe; @@ -1042,6 +1262,9 @@ pub(crate) async fn run( app.notification_service.protocol(), app.notification_service.config().condition, ); + // Deduplicate by category: general terminal warnings take priority + // over notification-specific ones (e.g. DcsPassthrough can fire from + // both sources when allow-passthrough is off). let mut seen = std::collections::HashSet::new(); for w in &warnings { seen.insert(w.category); @@ -1055,7 +1278,17 @@ pub(crate) async fn run( if !all_warnings.is_empty() { tracing::info!("Collected {} startup warnings", all_warnings.len()); } + // WezTerm without the Kitty keyboard protocol breaks local input + // (Shift+Enter can't insert newlines), so its banner is surfaced + // directly (no SSH gate) and first — see `assemble_startup_warnings`. + // `xtversion::detected()` is structurally `None` here (the probe is + // only sent further down, right before the input reader thread is + // spawned), so this banner covers env-detected WezTerm; the SSH shape + // surfaces in /doctor once the async reply has landed. let wezterm_warning = crate::diagnostics::wezterm_kitty_keyboard_warning(&snapshot); + // Wayland no-data-control is surfaced without the SSH gate of + // `summarize_warnings` — the broken shape is local (see + // `assemble_startup_warnings`). let wayland_clipboard_warning = all_warnings .iter() .find(|w| w.category == crate::diagnostics::WarningCategory::WaylandNoDataControl); @@ -1070,7 +1303,11 @@ pub(crate) async fn run( .collect(), ); } + + // Apply initial config (may come from existing ~/.grok/pager.toml). let mut initial_config = config_watcher.current().clone(); + // The cache holds the USER compact value; the render value is derived + // (auto-compact while the startup terminal is short). initial_config.prompt.compact = crate::views::agent::effective_compact( crate::appearance::cache::load(), app.last_known_terminal_rows, @@ -1080,7 +1317,14 @@ pub(crate) async fn run( let tick_interval = initial_config.animation.tick_interval(); crate::appearance::set_tab_width(initial_config.scrollback.display.tab_width); app.set_appearance(initial_config); + + // Seed app state from disk once at the I/O boundary so dispatch + // stays sans-IO. app.current_ui = load_initial_ui_config(); + // Field-tolerant: a whole-`UiConfig` default (malformed unrelated `[ui]` + // field) must not wipe a valid `show_timeline` or leave appearance / + // cache / `current_ui` disagreeing — `/timeline` and the rail all read + // the same canonical value after this sync + `prime` below. let show_timeline = crate::appearance::cache::load_show_timeline(); app.current_ui.show_timeline = Some(show_timeline); if app.appearance.show_timeline != show_timeline { @@ -1088,13 +1332,21 @@ pub(crate) async fn run( config.show_timeline = show_timeline; app.set_appearance(config); } + // Single-key load so a malformed unrelated `[ui]` field cannot wipe this. let page_flip_on_send = crate::appearance::cache::load_page_flip_on_send(); app.current_ui.page_flip_on_send = Some(page_flip_on_send); + // Disk load replaces `current_ui`. Assign one policy-clamped resolved + // launch mode unconditionally (CLI > TOML > remote > Ask) so disk Auto + // cannot win over `--permission-mode ask`, and a policy-clamped remote + // AlwaysApprove cannot leave the UI claiming AlwaysApprove while + // enforcement is Ask. let display_mode: &'static str = if launch_auto { "auto" } else if launch_yolo.yolo { "always-approve" } else if let Some(cli) = args.permission_mode_flag.as_deref() { + // CLI always-approve/auto that did not become launch_yolo/launch_auto + // (policy pin / gate) display as Ask. xai_grok_shell::util::config::clamped_display_permission_mode( xai_grok_shell::util::config::parse_permission_mode_canonical(cli), ) @@ -1106,20 +1358,38 @@ pub(crate) async fn run( }; app.current_ui.permission_mode = Some(display_mode.to_string()); super::dispatch::downgrade_displayed_auto_if_gated(&mut app); + // Seed `/auto` feature-gate visibility from the resolved gate (so `/auto` + // is offered on the welcome prompt when available). app.sync_permission_mode_slash_gate(); + // Settings UI language (`[ui].voice_stt_language`) overrides `[voice].language` + // when set. Store the preference (including client-only `auto`); the voice + // crate resolves the wire code at STT connect. When unset, keep whatever + // `from_config_table` loaded (default `en`, or an explicit `[voice].language`). + // Must run after `load_initial_ui_config()` hydrates `current_ui` from disk. if let Some(ref pref) = app.current_ui.voice_stt_language { app.voice_config.language = crate::settings::canonical_voice_stt_language(Some(pref)).to_string(); } + // Seed the Voice shortcut gate's process-global mirror for key-routing and + // view code without an `AppView`; the chord intercept reads `current_ui` + // live and the settings setter updates both. crate::app::VOICE_KEYBIND_ENABLED.store( app.current_ui.voice_keybind_enabled.unwrap_or(true), std::sync::atomic::Ordering::Release, ); + // Resolve the per-tip contextual hints now that `current_ui` is hydrated and + // propagate the prompt-relevant tips to any agents built at startup. New + // agents adopt the gates at creation; settings toggles re-apply at runtime. let resolved_hints = xai_grok_shell::util::config::resolve_contextual_hints( &app.current_ui.contextual_hints, app.remote_contextual_hints.as_ref(), ); app.apply_contextual_hints(resolved_hints); + + // Opt-in mouse-reporting toggle shortcut (Ctrl+R on scrollback). Off unless + // explicitly enabled. Resolved in shell config (env override > effective + // config > the parsed `UiConfig` field) so a partial `UiConfig` deserialize + // failure cannot silently drop it. let mouse_toggle = xai_grok_shell::util::config::resolve_mouse_reporting_toggle( effective_config.as_ref(), &app.current_ui, @@ -1128,6 +1398,8 @@ pub(crate) async fn run( term_state.screen_mode, mouse_toggle.value, ); + // Cache the resolved flag so the `/toggle-mouse-reporting` slash command can + // gate its visibility/execution without re-reading config on every keystroke. crate::app::MOUSE_REPORTING_TOGGLE_ENABLED .store(mouse_toggle.value, std::sync::atomic::Ordering::Release); let action_registered = app @@ -1152,29 +1424,71 @@ pub(crate) async fn run( app.show_tips = config_session_bools.show_tips; app.auto_update = config_session_bools.auto_update; app.ask_user_question_timeout_enabled = config_session_bools.ask_user_question_timeout_enabled; + // Prime thread-local caches so first render doesn't hit disk. crate::appearance::cache::prime(&app.current_ui); + // Re-derive the render-value compact flag from the hydrated `current_ui`: + // the seed above used the pre-hydration disk read, which layered/remote + // config can contradict — the canonical single-writer corrects it (and + // fans out to any startup agents) before the first draw. app.apply_effective_compact(); + + // Apply the scroll settings from the caches (seeded by `prime` above; + // GROK_SCROLL_SPEED/_MODE/_LINES + GROK_INVERT_SCROLL env overrides + // apply on first load). app.scroll_config = crate::input::mouse::ScrollConfig::from_settings(); + + // Fire-and-forget XTVERSION query; must sit immediately before the input + // reader thread is spawned so no earlier stdin consumer eats the reply. + // DA2 shares that constraint but runs earlier, in `init_terminal`, so its + // version is already resolved when the startup telemetry above is emitted. crate::terminal::xtversion::probe_at_startup(); + + // Read terminal events on a dedicated thread and forward them over an mpsc + // channel. The main `select!` consumes via `input_rx.recv()`, which is + // cancellation-safe: when another arm wins, the recv future is dropped and + // re-created without losing the wakeup. Polling crossterm's `EventStream` + // directly in the select is NOT safe -- dropping its `next()` future + // mid-poll (a losing arm) strands its background waker (crossterm #936), so + // input on an idle screen was not serviced until an unrelated arm happened + // to re-poll (every ~20s via recap_poll). The always-on tracing_rx tick + // used to mask this by re-polling ~30Hz; this removes that dependency. let (input_tx, mut input_rx) = tokio::sync::mpsc::unbounded_channel::(); + // Set true around tty handoffs (e.g. $EDITOR) so the reader stops touching + // stdin and the inheriting child process keeps every keystroke. The handoff + // does not proceed until `reader_parked` acknowledges this pause. let input_paused = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); let reader_paused = input_paused.clone(); + // Set by the reader once it has parked (stopped calling crossterm) so the + // $EDITOR handoff can wait for it: poll/read share one global lock, so the + // main-thread drain must be the sole crossterm caller. let reader_parked = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); let reader_parked_thread = reader_parked.clone(); std::thread::spawn(move || { use std::sync::atomic::Ordering; + // Short enough that a pause / receiver-drop is observed promptly, long + // enough to keep the thread parked when idle. A `poll()` timeout here + // does NOT wake the main loop -- only a successful `send` does -- so the + // idle event loop still parks (no reintroduced metronome tick). const POLL_TIMEOUT: Duration = Duration::from_millis(100); let mut consecutive_event_errors: u32 = 0; loop { + // Shutdown observed within one poll cycle in every state (idle or + // paused); the send() break below covers close-while-sending. if input_tx.is_closed() { break; } + // While a tty handoff owns stdin, do not read(): the child (e.g. the + // editor) must keep its bytes. Re-check soon without touching stdin. if reader_paused.load(Ordering::Acquire) { + // Signal the handoff that the reader is no longer in crossterm. reader_parked_thread.store(true, Ordering::Release); std::thread::sleep(POLL_TIMEOUT); continue; } + // Active path: this thread owns crossterm again this iteration. reader_parked_thread.store(false, Ordering::Release); + // poll()+read() (not a bare blocking read) so the pause flag and a + // dropped receiver are observed within POLL_TIMEOUT. let event = match crossterm::event::poll(POLL_TIMEOUT) { Ok(true) => crossterm::event::read(), Ok(false) => continue, @@ -1185,10 +1499,13 @@ pub(crate) async fn run( consecutive_event_errors = 0; let timed = TimedInputEvent::now(ev); if input_tx.send(timed).is_err() { - break; + break; // event loop has shut down } } Err(e) => { + // VTE terminals / SSH PTYs can emit garbage that crossterm's + // parser rejects; skip transient errors rather than kill the + // TUI (ratatui#1275), bailing only if they never stop. consecutive_event_errors += 1; if consecutive_event_errors >= 50 { tracing::error!( @@ -1208,41 +1525,85 @@ pub(crate) async fn run( let mut tasks: JoinSet = JoinSet::new(); let (progress_tx, mut progress_rx) = tokio::sync::mpsc::unbounded_channel::(); + + // Voice STT pipeline is started lazily on first successful `/voice` (see + // `VoiceState::ColdStart`), not at launch — avoids background work for users + // who never enable voice mode. `AUDIO_SUPPORTED` reflects whether mic + // capture is compiled in: true for production CLI builds on macOS/Windows + // (cpal) and Linux (subprocess recorder), false for Bazel builds (no + // capture in the test sandbox). let mut voice_rx = None::>; let voice_auth_factory = connection.auth_manager.clone(); + + // Animation tick: only scheduled when there are running entries. let mut tick_interval = tick_interval; let mut animation_tick_at: Option = None; + + // Whether the extra Kitty keyboard layer (WASD release events) is + // currently pushed for the /gboom game. Synced to `gboom_active` each + // 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 + // without waiting a full interval. const ROSTER_POLL_INTERVAL: Duration = Duration::from_secs(1); let mut roster_poll_at: Option = Some(Instant::now()); + + // Pre-generate the automatic "return-from-away" recap while the terminal is + // unfocused, so it's already in the scrollback (instant) when the user + // returns. The arm is a cheap no-op while focused / not-yet-eligible; the + // heavy lifting (the model call) only fires once per away period via + // `should_pregenerate_away_recap`. const RECAP_POLL_INTERVAL: Duration = Duration::from_secs(20); let mut recap_poll_at: Option = Some(Instant::now() + RECAP_POLL_INTERVAL); + + // Seed the folder-trust verdict BEFORE the first render and before any + // session is created (no repo-local MCP/LSP/hooks/plugins have loaded yet). + // Feature-off (kill-switch / opt-out / local build) resolves `Trusted`, so + // this stays `TrustState::Done`. seed_trust_state(&mut app, remote_settings.as_ref()); + let mut presenter = Presenter::new(); + // A timed-out handoff stays queued but cannot synchronously retry until + // this deadline fires. Feedback is one-shot per editor/pager request, even + // across multiple deferred attempts. let mut suspend_retry_after: Option = None; let mut suspend_wait_reports = SuspendWaitReports::default(); + + // Initial render presenter.request_presentation(&mut app, terminal, false); + + // status only; shell auto-syncs post-auth if matches!(app.auth_state, AuthState::Done) { let effs = dispatch::dispatch(Action::RequestBundleStatus, &mut app); 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)); @@ -1251,11 +1612,16 @@ pub(crate) async fn run( gate_poll_at = Some(Instant::now() + GATE_POLL_INTERVAL); } } + if !post_render_effects.is_empty() && process_effects(post_render_effects, &mut tasks, &mut app, &progress_tx) { return Ok(make_run_result(&app)); } + + // Session startup from pre-materialized CLI intent. + // These actions are dispatched UNCONDITIONALLY: the session-creating + // chokepoints self-gate when auth + folder trust is closed. use crate::app::session_startup::MaterializedStartup; let startup_action = match &materialized { MaterializedStartup::Resume { @@ -1268,6 +1634,8 @@ pub(crate) async fn run( restore_code = ?app.restore_code, "RESTORE_CODE_DEBUG: worktree+resume path taken" ); + // Materialization-time provenance for the worktree failure hint; + // the effect matches it against the exact deferred target. app.resume_local_miss = deferred_local_miss.then(|| session_id.clone()); Some(Action::NewWorktreeSession { load_session_id: Some(session_id.clone()), @@ -1275,12 +1643,21 @@ pub(crate) async fn run( git_ref: args.worktree_ref.clone(), }) } - MaterializedStartup::Resume { session_id, .. } => Some(Action::LoadSession( - session_id.clone(), - session_cwd.clone(), - false, - )), + MaterializedStartup::Resume { session_id, .. } => { + // CLI resume has no roster entry: `chat_kind` on LoadSession is the + // conversation-entry bit only (false here). Process-wide `--chat` + // still stamps kind=chat via SessionFlags.chat_mode in the load + // effect; local Build disk rows are refused in dispatch / startup. + Some(Action::LoadSession( + session_id.clone(), + session_cwd.clone(), + false, + )) + } MaterializedStartup::NewWithId { session_id } if args.worktree.is_some() => { + // Stash preferred id; `dispatch_new_worktree_session` consumes it and + // passes through `CreateWorktreeSession.preferred_session_id` so the + // worktree + ACP session use the CLI-chosen id (not an auto `pager-*`). app.deferred_startup.preferred_session_id = Some(session_id.clone()); Some(Action::NewWorktreeSession { load_session_id: None, @@ -1310,6 +1687,7 @@ pub(crate) async fn run( } MaterializedStartup::NewAuto => None, }; + if let Some(action) = startup_action { let effs = dispatch::dispatch(action, &mut app); if process_effects(effs, &mut tasks, &mut app, &progress_tx) { @@ -1317,6 +1695,7 @@ pub(crate) async fn run( } presenter.request_presentation(&mut app, terminal, false); } else if args.worktree.is_some() { + // --worktree only: create worktree + new session. let effs = dispatch::dispatch( Action::NewWorktreeSession { load_session_id: None, @@ -1330,6 +1709,13 @@ pub(crate) async fn run( } presenter.request_presentation(&mut app, terminal, false); } + + // Initial prompt from the CLI positional (`grok "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`. if let Some(initial_prompt) = args.initial_prompt() { if !app.session_startup_allowed() { app.deferred_startup.prompt = Some(initial_prompt.to_string()); @@ -1341,7 +1727,12 @@ pub(crate) async fn run( presenter.request_presentation(&mut app, terminal, false); } } + + // `grok dashboard` startup: open the dashboard view immediately. The + // CLI subcommand wrote a `GROK_OPEN_DASHBOARD_AT_STARTUP=1` env var + // so we don't have to thread a flag through every arg struct. if std::env::var("GROK_OPEN_DASHBOARD_AT_STARTUP").as_deref() == Ok("1") { + // SAFETY: we are pre-multithreaded init for this app loop. unsafe { std::env::remove_var("GROK_OPEN_DASHBOARD_AT_STARTUP") }; if app.session_startup_allowed() { let effs = dispatch::dispatch(Action::OpenDashboard, &mut app); @@ -1350,45 +1741,111 @@ pub(crate) async fn run( } presenter.request_presentation(&mut app, terminal, false); } else { + // Not signed in yet — the env var is already consumed, so + // without a stash the request would be silently dropped and + // the post-login flow would land on the welcome screen. + // Defer to the `AuthComplete` handler (mirrors + // the deferred session/prompt owner). app.deferred_startup.open_dashboard = true; } } + + // Minimal (scrollback-native) mode has no welcome screen: the live region + // only renders for an Agent view. If nothing above already started a + // session (no resume / initial prompt / worktree / dashboard), open an + // 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 app.session_startup_allowed() { + // Already authenticated + trusted: open the empty session now so the + // user lands directly at the prompt. let effs = dispatch::dispatch(Action::NewSession, &mut app); if process_effects(effs, &mut tasks, &mut app, &progress_tx) { return Ok(make_run_result(&app)); } presenter.request_presentation(&mut app, terminal, false); } else { + // Sign-in (or folder-trust) still pending: minimal renders the + // device / external sign-in flow in its live region. Defer the + // empty-session creation so the post-auth (or post-trust) drain + // (`drain_startup_actions`) opens it — otherwise minimal would + // authenticate but never create a session, stranding the user on the + // sign-in screen. app.deferred_startup.new_session = true; } } + + // Startup intents are now fully classified; only an untouched welcome can nudge. if let Some(effect) = app.begin_foreign_resume_detection() && process_effects(vec![effect], &mut tasks, &mut app, &progress_tx) { return Ok(make_run_result(&app)); } + + // Schedule the first animation tick so live updates start immediately + // (without waiting for user input). schedule_tick(&mut animation_tick_at, &app, tick_interval); + + // Resize debounce: during continuous terminal drags, dozens of resize + // events fire per second. Each would trigger a full layout rebuild of all + // entries (the most expensive per-frame operation). Instead of drawing on + // every resize, we schedule a single deferred draw after the size stabilizes. const RESIZE_DEBOUNCE: Duration = Duration::from_millis(16); let mut resize_debounce_at: Option = None; + + // Cadences resolved once above (env > auto > 16ms). AppView/Default stays hermetic. app.scroll_state.set_redraw_cadence(scroll_cadence); + // ACP batch bound: large enough to keep the hundreds-buffered streaming + // case batched (draws stay cadence-throttled regardless), small enough that + // loop-top work (suspends, deadline re-derivation) never waits on an + // unbounded drain during a token firehose. const ACP_DRAIN_BATCH_MAX: usize = 32; + let mut reconnect_reinit: Option = None; let mut reconnect_abort_handle: Option = None; + // Highest `Connected` generation already handled. Starts at 0 — the + // initial pre-reconnect watch value — so startup never triggers a reload; + // any greater generation is a reconnect, even when the intermediate + // `Reconnecting` state was coalesced away by the watch channel. let mut last_leader_generation: u64 = 0; + + // Persistent CSI fragment filter — carries parsing state across + // drain_and_process calls so a mouse report split across batches is still + // caught; a focus report is only swallowed when its `\e` and `[I`/`[O` + // land in the same batch. let mut csi_filter = super::csi_filter::CsiFragmentFilter::new(); + + // Swallows the fire-and-forget XTVERSION reply whenever it arrives; + // armed only when the startup query is still unanswered. let mut xt_filter = super::xt_filter::XtversionFilter::new(); + + // Background update check: resolves when the spawned update task + // determines whether a newer version is available. let mut bg_update_rx = bg_update_rx; + + // `app::run` publishes the resolved theme into `theme_cache::CURRENT` + // before `init_terminal` so `apply_cursor_color()` sees it. Pin the + // invariant so a future refactor that drops the `theme_cache::set` call + // fails loudly in debug builds rather than silently regressing the + // initial cursor color. debug_assert_eq!(term_state.initial_theme, theme_cache::current_kind()); let mut appearance_watcher = SystemAppearanceWatcher::start_if_auto(theme_cache::is_auto_mode()); + + // Registered so the signal handler can request a graceful quit; see signal_handler. let quit_notify = std::sync::Arc::new(tokio::sync::Notify::new()); crate::app::signal_handler::set_quit_notify(quit_notify.clone()); + loop { + // Pending $EDITOR / $PAGER suspends first: they can be armed by ANY + // arm of the select below (input, ticks — e.g. minimal's incremental + // /transcript build finishing inside a tick draw — tasks, ACP), so + // consuming them here keeps the handoff immediate instead of waiting + // for the next unrelated event. run_pending_suspends( &mut app, terminal, @@ -1399,6 +1856,10 @@ pub(crate) async fn run( &mut suspend_retry_after, &mut suspend_wait_reports, )?; + + // Lazy voice pipeline: only after `/voice` or Ctrl+Space while gates + // allow. Consume the queued cold-start, carrying its hold-ownership and + // bound target forward into the live recording it spawns. if let VoiceState::ColdStart { hold, target } = app.voice_state { if app.voice_cmd_tx.is_none() && app.voice_can_start_pipeline() { let voice_auth = crate::voice::build_voice_auth(voice_auth_factory.clone()); @@ -1415,6 +1876,13 @@ pub(crate) async fn run( app.voice_cmd_tx = Some(cmd_tx); voice_rx = Some(event_rx); tracing::info!("voice pipeline started (/voice or Ctrl+Space)"); + // The spawn is async, so begin capture now the pipeline is live + // — but only if the user is still on a surface that can receive + // dictation (an agent prompt or the dashboard dispatch input). + // This runs at loop-top before any new input, so the surface + // normally can't have changed since the keypress; the else-arm + // is defensive cleanup so voice mode can't stay armed without + // capture ever starting. if matches!( app.active_view, ActiveView::Agent(_) | ActiveView::AgentDashboard @@ -1429,56 +1897,100 @@ pub(crate) async fn run( app.voice_ui_active = false; app.show_toast("Voice could not start. Restart Grok."); } else { + // Defensive: a queued start with the pipeline already up (which + // shouldn't occur) — drop it so we don't re-enter every tick. app.voice_state = VoiceState::Idle; } + // The lazy spawn runs at loop-top, after the key/slash arm already + // drew (with capture still off). Render now so the recording banner + // appears immediately instead of waiting for the next input or + // network event to wake the select! loop. presenter.request_presentation(&mut app, terminal, false); } + + // Stop voice if the user has left the recording session (see method). app.enforce_voice_session_bound(); + + // Keep the /gboom keyboard layer in sync with whether the game is + // open, so WASD emit releases while it runs and the layer is popped + // on every close path (Esc, game-over dismiss, session switch). let want_gboom_keyboard = app.gboom_active(); if want_gboom_keyboard { if !gboom_keyboard_pushed { super::push_gboom_keyboard_flags(); gboom_keyboard_pushed = true; } + // Only the active game receives release events; any other open + // game must drop its latched holds, or it resumes walking with + // no key down when reopened after a tab/view switch. app.gboom_release_backgrounded_games(); } else if gboom_keyboard_pushed { super::pop_gboom_keyboard_flags(); gboom_keyboard_pushed = false; + // No game is the active input target now (switched to a non-game + // view); clear every game's holds for the same reason. app.gboom_release_all_games(); } + + // Re-arm the dashboard roster poll when the dashboard is open but the + // poll has gone dormant — i.e. the dashboard was just opened. The poll + // arm leaves `roster_poll_at = None` only when it fired with the + // dashboard closed, so this fires an immediate refresh exactly on the + // closed→open transition rather than every iteration. Applies in both + // modes: leader mode polls the live roster, non-leader mode polls the + // local on-disk idle-session list. if roster_poll_at.is_none() && matches!(app.active_view, ActiveView::AgentDashboard) { 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 { Some(at) => sleep_until(at).await, None => std::future::pending().await, } }; + + // Dedicated scroll clock, derived fresh each iteration — a pure + // function of scroll state, so no arm can forget to reschedule it. + // Armed only while a wheel/trackpad stream is active, at the state + // machine's own deadline (16ms cadence flushes while lines are + // pending, the 80ms stream-gap finalize otherwise): scroll pacing + // must never ride the slower animation fps, which turned residual + // flushes into visible jumps. let scroll_tick_at = { let now = Instant::now(); app.scroll_state .scroll_clock_deadline(now.into_std()) .map(|delay| now + delay) }; + + // Future that sleeps until the scroll deadline, or waits forever. let scroll_tick = async { match scroll_tick_at { Some(at) => sleep_until(at).await, None => std::future::pending().await, } }; + + // Future that sleeps until the resize debounce fires, or waits forever. let resize_debounce = async { match resize_debounce_at { Some(at) => sleep_until(at).await, None => std::future::pending().await, } }; + + // Future that sleeps until a throttled draw fires, or waits forever. let deferred_draw_at = presenter.draw_scheduled_at; let deferred_draw = async move { match deferred_draw_at { @@ -1486,6 +1998,8 @@ pub(crate) async fn run( None => std::future::pending().await, } }; + + // Wake a deferred suspend retry without requiring unrelated input. let suspend_retry_at = if app.pending_editor.is_some() || app.pending_pager_path.is_some() { suspend_retry_after } else { @@ -1497,36 +2011,42 @@ pub(crate) async fn run( None => std::future::pending().await, } }; + 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, None => std::future::pending().await, } }; + let recap_poll = async { match recap_poll_at { Some(at) => sleep_until(at).await, None => std::future::pending().await, } }; + tokio::select! { biased; @@ -2284,11 +2804,15 @@ pub(crate) async fn run( } } } + presenter.present_if_dirty(&mut app, terminal); } + app.notification_service.shutdown(); + Ok(make_run_result(&app)) } + /// Load `UiConfig` from the shell's layered config at startup. /// Falls back to `UiConfig::default()` on any failure. pub(crate) fn load_initial_ui_config() -> xai_grok_shell::agent::config::UiConfig { @@ -2301,6 +2825,7 @@ pub(crate) fn load_initial_ui_config() -> xai_grok_shell::agent::config::UiConfi }; ui_value.try_into::().unwrap_or_default() } + /// Config `Option` mirrors seeded once at startup. `None` = no /// TOML override; the modal falls back to the per-setting default. #[derive(Default)] @@ -2309,6 +2834,7 @@ struct InitialConfigSessionBools { auto_update: Option, ask_user_question_timeout_enabled: Option, } + fn load_initial_config_session_bools() -> InitialConfigSessionBools { let Ok(root) = xai_grok_shell::config::load_effective_config() else { return InitialConfigSessionBools::default(); @@ -2324,6 +2850,7 @@ fn load_initial_config_session_bools() -> InitialConfigSessionBools { .and_then(|v| v.as_bool()), } } + /// Whether to pre-generate the automatic "return-from-away" recap right now. /// /// True only when the terminal has been unfocused past the recap threshold @@ -2346,6 +2873,7 @@ fn apply_session_recap_available(app: &mut AppView, available: bool) { dashboard.set_recap_visible(available); } } + fn should_pregenerate_away_recap(app: &AppView) -> bool { if !(app.session_recap_available && app.notification_service.focus_tracker.recap_due() @@ -2364,15 +2892,20 @@ fn should_pregenerate_away_recap(app: &AppView) -> bool { && !agent.session.has_running_bg_tasks() }) } + /// Schedule the next animation tick when demanded and none is pending. fn schedule_tick(tick_at: &mut Option, app: &AppView, interval: Duration) { if tick_at.is_none() { let interval = match app.tick_demand() { crate::app::app_view::TickDemand::None => return, + // A view can request a faster cadence than the configured + // animation fps (e.g. the /gboom easter egg targets ~30 fps). crate::app::app_view::TickDemand::Fast => match app.tick_interval_ceiling() { Some(ceiling) => interval.min(ceiling), None => interval, }, + // Only low-frequency work (welcome shimmer, Cmd link poll): + // don't spin the full 30fps loop for it. crate::app::app_view::TickDemand::Slow => { interval.max(crate::app::app_view::SLOW_TICK_INTERVAL) } @@ -2380,6 +2913,7 @@ fn schedule_tick(tick_at: &mut Option, app: &AppView, interval: Duratio *tick_at = Some(Instant::now() + interval); } } + /// Sync `appearance_watcher` with the current `AUTO_MODE` flag. /// Starts or stops the watcher as needed; no-op when consistent. fn sync_appearance_watcher(watcher: &mut Option) { @@ -2388,6 +2922,7 @@ fn sync_appearance_watcher(watcher: &mut Option) { *watcher = SystemAppearanceWatcher::start_if_auto(should_auto); } } + /// Build [`ExitInfo`] from the active agent's session (if any). /// /// Sole construction site of [`super::ExitSummary`]: fullscreen quits only @@ -2426,6 +2961,7 @@ fn make_run_result(app: &AppView) -> RunResult { relaunch: app.relaunch.clone(), } } + /// Result of draining and processing terminal events. struct DrainResult { /// Whether any event produced a visual change requiring a draw. @@ -2440,14 +2976,17 @@ struct DrainResult { /// refocus in editor/multiplexer contexts to heal out-of-band stranded rows. force_repaint: bool, } + struct RoutedInputEvent { event: Event, arrived_at: std::time::Instant, paste_provenance: PasteProvenance, } + fn tty_suspend_armed(app: &AppView) -> bool { app.pending_editor.is_some() || app.pending_pager_path.is_some() } + fn normalize_input_event(timed: TimedInputEvent) -> RoutedInputEvent { let TimedInputEvent { event, arrived_at } = timed; #[cfg(target_os = "linux")] @@ -2476,6 +3015,7 @@ fn normalize_input_event(timed: TimedInputEvent) -> RoutedInputEvent { paste_provenance: PasteProvenance::Terminal, } } + /// Process a terminal event, then drain any buffered events before returning. /// /// Crossterm buffers input events while the app is drawing. Without draining, @@ -2500,19 +3040,34 @@ async fn drain_and_process( let mut had_resize = false; let mut had_non_resize_change = false; let mut force_repaint = false; + + // Collect all immediately-available events for paste coalescing. let mut raw_events = vec![first]; drain_immediate(&mut raw_events, input_rx); + + // XTVERSION reply removal must precede paste coalescing so reply chars + // are never folded into a synthetic Paste. if xt_filter.armed() { raw_events = super::xt_filter::filter_with_fragment_wait(xt_filter, raw_events, input_rx).await; } + + // On terminals without bracketed paste, try to capture more events + // that may still be in transit from the input reader thread. if should_extend_for_paste(&raw_events) && detect_paste(&mut raw_events, input_rx).await { collect_remaining_paste(&mut raw_events, input_rx).await; + // The paste extension pulled more events off the channel without + // running them through the still-armed filter — a late or split + // XTVERSION reply could otherwise be folded into the paste. if xt_filter.armed() { raw_events = super::xt_filter::filter_with_fragment_wait(xt_filter, raw_events, input_rx).await; } } + + // The /gboom game tracks keys by press → release, so it needs the + // release events that `coalesce_rapid_keys` strips (and it never + // pastes). Skip coalescing while it owns input. let coalesced = if app.gboom_active() { raw_events } else { @@ -2523,28 +3078,45 @@ async fn drain_and_process( .into_iter() .map(normalize_input_event) .collect::>(); + let suspend_armed_after_event = std::cell::Cell::new(false); let mut handle_one = |routed: &RoutedInputEvent| -> bool { let ev = &routed.event; match ev { Event::FocusGained => { + // Force a full repaint on refocus to heal out-of-band stranded rows. + // Sets needs_draw (not had_non_resize_change); the draw site honors force_repaint + // ahead of the resize debounce, clearing even a coalesced same-size resize. if crate::terminal::terminal_context().repaints_pane_out_of_band() { force_repaint = true; needs_draw = true; } + // Capture recap eligibility BEFORE on_focus_gained() clears the + // away timer. Auto recap requires the shell rollout flag plus + // the notifications opt-in; manual `/recap` only needs the flag. let recap_due = app.session_recap_available && app.notification_service.focus_tracker.recap_due() && app.notification_service.config().session_recap; app.notification_service.focus_tracker.on_focus_gained(); + // Pre-warm AppKit's lazy dlopen off the UI thread (once) so the + // first changeCount poll after returning is just the cheap + // metadata read and never stalls a frame on the framework load. + // FocusGained is itself an active loop iteration, so the + // opportunistic poll (driven after drain_and_process) does the + // actual clipboard check — no debounce, no timer, and + // `needs_animation` is never kept hot for it. if app.contextual_hints.image_input && crate::clipboard::clipboard_image_probe_supported() { 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) => { if let Some(agent) = app.agents.get_mut(&id) @@ -2554,6 +3126,12 @@ async fn drain_and_process( needs_draw = true; had_non_resize_change = true; } + + // Automatic "where was I" recap: the user just returned + // after being away long enough. Only when the session is + // idle and not blocked by a modal or pending question. + // Compute eligibility into a bool first so the immutable + // agent borrow is dropped before dispatch (&mut app). let eligible = app.agents.get(&id).is_some_and(|agent| { agent.session.state.is_idle() && agent.active_modal.is_none() @@ -2581,12 +3159,17 @@ async fn drain_and_process( had_non_resize_change = true; } } + // The dashboard manages its own input/overview focus + // (`list_focused`); refocusing the terminal must not + // override the user's choice (e.g. vim overview focus). ActiveView::AgentDashboard => {} } return false; } Event::FocusLost => { app.notification_service.focus_tracker.on_focus_lost(); + // The /gboom game latches held keys until their release; a + // release can be lost while unfocused, so stop all movement. if app.gboom_active() { app.gboom_release_all_games(); needs_draw = true; @@ -2595,6 +3178,15 @@ async fn drain_and_process( } _ => {} } + // Voice capture chord (Ctrl+Space or F8), handled here before normal + // routing so the release reaches us and the key never lands as text. + // Hold-to-talk where releases are reported (press records, release + // stops), else tap toggle. A release is only ours when a hold session + // owns it, so a bare Space release (Ctrl lifted first) stops + // hold-to-talk without eating every Space release during normal typing. + // `[ui].voice_keybind_enabled` (read live, like `voice_capture_mode`) + // silences chord presses without touching `/voice` — see + // `voice_chord_claims_event` for the exact press/release/hold gating. if let Event::Key(ke) = ev && app.voice_mode_enabled && xai_grok_voice::AUDIO_SUPPORTED @@ -2605,6 +3197,8 @@ async fn drain_and_process( app.voice_hold_owned(), ) { + // Hold-to-talk only when selected AND the terminal reports key + // releases (Kitty protocol); otherwise fall back to a tap toggle. let hold_mode = crate::settings::canonical_voice_capture_mode( app.current_ui.voice_capture_mode.as_deref(), ) == "hold"; @@ -2640,6 +3234,9 @@ async fn drain_and_process( had_non_resize_change = true; } InputOutcome::ActionThenForward(action) => { + // Dispatch the action (e.g. create session), then re-process + // the same event through the now-active view so the input + // (character, paste) lands in the session's prompt. let effs = dispatch::dispatch(action, app); if process_effects(effs, tasks, app, progress_tx) { return true; @@ -2658,6 +3255,8 @@ async fn drain_and_process( had_non_resize_change = true; } InputOutcome::ActionPair(first, second) => { + // Dispatch both in order; first must fully resolve + // before second (e.g. revert preview then open reset). let effs = dispatch::dispatch(first, app); if process_effects(effs, tasks, app, progress_tx) { return true; @@ -2677,6 +3276,7 @@ async fn drain_and_process( had_non_resize_change = true; } } + // AppView converts ArmPending → Changed; defensive if one slips through. InputOutcome::ArmPending { .. } => { needs_draw = true; had_non_resize_change = true; @@ -2686,6 +3286,7 @@ async fn drain_and_process( suspend_armed_after_event.set(tty_suspend_armed(app)); false }; + for routed in &coalesced { if handle_one(routed) { return DrainResult { @@ -2695,10 +3296,12 @@ async fn drain_and_process( force_repaint: false, }; } + // Hand off to the TTY-taking child before later buffered events mutate UI state. if suspend_armed_after_event.get() { break; } } + DrainResult { needs_draw, should_quit: false, @@ -2706,19 +3309,26 @@ async fn drain_and_process( force_repaint, } } + +// ── Paste coalescing for terminals without bracketed paste ─────────── + /// Timeout for the first extension round (detection). If no event /// arrives within this window the batch was a normal keystroke. const PASTE_DETECT_TIMEOUT: Duration = Duration::from_millis(2); + /// Timeout for subsequent rounds once paste has been detected. const PASTE_CONTINUE_TIMEOUT: Duration = Duration::from_millis(10); + /// Safety cap on events accumulated in one extension pass. const PASTE_EXTEND_MAX_EVENTS: usize = 5_000; + /// Returns `true` when the batch contains pasteable key events but no /// `Event::Paste` (i.e. bracketed paste is not handling it). fn should_extend_for_paste(events: &[TimedInputEvent]) -> bool { !events.iter().any(|e| matches!(e.event, Event::Paste(_))) && events.iter().any(|e| is_pasteable_key_event(&e.event)) } + /// Wait [`PASTE_DETECT_TIMEOUT`] for a follow-up event. Returns `true` /// if a **pasteable key event** arrives within the window. Non-key events /// (mouse, focus, releases) are collected but do not count as paste evidence. @@ -2738,6 +3348,7 @@ async fn detect_paste( _ => false, } } + /// Collect remaining paste events using [`PASTE_CONTINUE_TIMEOUT`]. /// Only pasteable key events extend the timeout; non-key events are /// collected but do not keep the loop alive. @@ -2767,6 +3378,7 @@ async fn collect_remaining_paste( } } } + /// Non-blocking drain of all immediately available events. pub(super) fn drain_immediate( batch: &mut Vec, @@ -2776,13 +3388,16 @@ pub(super) fn drain_immediate( batch.push(ev); } } + /// Minimum key events in a run to trigger paste coalescing. const PASTE_COALESCE_THRESHOLD: usize = 3; + /// Minimum run length for the Windows path-shape coalesce branch. /// Covers the shortest realistic dropped image path (`C:\x.png`, /// `/a.png`) while leaving short typed prose alone. #[cfg(target_os = "windows")] const PATH_COALESCE_THRESHOLD: usize = 8; + /// Check if a terminal event is a pasteable key press — a character, /// Enter, or Tab with no control modifiers (Ctrl/Alt/Super). /// @@ -2802,6 +3417,7 @@ fn is_pasteable_key_event(ev: &Event) -> bool { _ => false, } } + /// Map a voice-chord key event to its action (pure, so it's unit-testable). /// /// Hold mode is press-to-record / release-to-stop, but only a hold-*owned* @@ -2820,7 +3436,7 @@ fn voice_chord_action( KeyEventKind::Press if !listening => Some(Action::EnableVoiceMode), KeyEventKind::Press if !hold_owned => Some(Action::VoiceToggle), KeyEventKind::Release => Some(Action::VoiceStop), - _ => None, + _ => None, // repeat while a hold is held, or press of a hold-owned session } } else if kind == KeyEventKind::Press { Some(Action::VoiceToggle) @@ -2828,6 +3444,7 @@ fn voice_chord_action( None } } + /// Whether the event-loop intercept claims a voice-chord key event (pure for /// unit tests). /// @@ -2843,6 +3460,7 @@ fn voice_chord_claims_event(kind: KeyEventKind, keybind_enabled: bool, hold_owne } kind != KeyEventKind::Release && keybind_enabled } + /// The voice-capture chord: **Ctrl+Space** or **F8**. A press needs the exact /// chord (matching the registry, so Shift+F8 / Ctrl+Alt+Space don't fire); a /// release matches the key alone (Space/F8), since on Kitty the Ctrl release can @@ -2857,6 +3475,7 @@ fn is_voice_chord(ke: &KeyEvent) -> bool { } } } + /// Coalesce runs of rapid key events into synthetic `Event::Paste` /// events. On terminals without bracketed paste, pasted text arrives /// as individual key events; Enter keys mid-run would otherwise @@ -2875,9 +3494,14 @@ fn is_voice_chord(ke: &KeyEvent) -> bool { /// /// No-op when bracketed paste already arrives as `Event::Paste`. fn coalesce_rapid_keys(events: Vec) -> Vec { + // Fast path: not enough events for coalescing to trigger. if events.len() < PASTE_COALESCE_THRESHOLD { return events; } + + // If Event::Paste fragments are mixed with key events (Windows + // Terminal can split a large bracketed paste across read boundaries), + // merge everything into a single Event::Paste. let (mut has_paste, mut has_keys) = (false, false); for e in &events { has_paste |= matches!(e.event, Event::Paste(_)); @@ -2890,6 +3514,9 @@ fn coalesce_rapid_keys(events: Vec) -> Vec { events }; } + + // Remove Release events — handlers ignore them and they'd break run + // detection. Exception: voice-chord releases (needed for hold-to-talk). let events: Vec = events .into_iter() .filter(|ev| { @@ -2897,8 +3524,10 @@ fn coalesce_rapid_keys(events: Vec) -> Vec { if ke.kind == KeyEventKind::Release && !is_voice_chord(ke)) }) .collect(); + let mut result = Vec::with_capacity(events.len()); let mut i = 0; + while i < events.len() { if is_pasteable_key_event(&events[i].event) { let run_start = i; @@ -2906,6 +3535,7 @@ fn coalesce_rapid_keys(events: Vec) -> Vec { let mut text = String::new(); let mut seen_enter = false; let mut has_char_after_enter = false; + while i < events.len() && is_pasteable_key_event(&events[i].event) { if let Event::Key(ke) = &events[i].event { match ke.code { @@ -2930,8 +3560,13 @@ fn coalesce_rapid_keys(events: Vec) -> Vec { } i += 1; } + let run_len = i - run_start; let multiline_paste = run_len >= PASTE_COALESCE_THRESHOLD && has_char_after_enter; + // Windows fallback for drag-drops that arrive as a key + // burst instead of a bracketed paste — reuse the drop + // classifier's anchor detector so the two layers can't + // drift on what counts as a path. #[cfg(target_os = "windows")] let path_shaped_drop = run_len >= PATH_COALESCE_THRESHOLD && crate::prompt_images::starts_with_drop_anchor(&text); @@ -2958,8 +3593,10 @@ fn coalesce_rapid_keys(events: Vec) -> Vec { i += 1; } } + result } + pub(super) fn is_bare_esc_press(ev: &Event) -> bool { matches!( ev, @@ -2968,6 +3605,7 @@ pub(super) fn is_bare_esc_press(ev: &Event) -> bool { && ke.modifiers == KeyModifiers::NONE ) } + /// Merge `Event::Paste` fragments and interleaved key events into a /// single `Event::Paste`. Non-paste, non-key events (Resize, Mouse, /// Focus) are preserved in order around the merged paste. @@ -2975,6 +3613,7 @@ fn merge_paste_fragments(events: Vec) -> Vec { let mut result = Vec::new(); let mut merged_text = String::new(); let mut merged_arrived_at = None; + for ev in events { match &ev.event { Event::Paste(text) => { @@ -2990,6 +3629,8 @@ fn merge_paste_fragments(events: Vec) -> Vec { _ => {} } } + // Non-pasteable keys (Ctrl+C, Backspace, arrows, Release + // events, etc.) are artifacts of paste fragmentation — drop. Event::Key(_) => {} _ => { if !merged_text.is_empty() { @@ -3004,22 +3645,83 @@ fn merge_paste_fragments(events: Vec) -> Vec { } } } + if !merged_text.is_empty() { result.push(TimedInputEvent { event: Event::Paste(merged_text), arrived_at: merged_arrived_at.expect("non-empty merged paste has an arrival time"), }); } + result } -/// Spawn effects into the task set. Returns `true` if the app should quit. -fn process_effects( - effs: Vec, - tasks: &mut JoinSet, - app: &mut AppView, - progress_tx: &tokio::sync::mpsc::UnboundedSender, + +/// True when this batch should consume the welcome local-workspace one-shot. +#[cfg(feature = "local-workspace")] +pub(crate) fn welcome_oneshot_applies_to_effects(effs: &[super::actions::Effect]) -> bool { + use super::actions::Effect; + effs.iter().any(|e| { + matches!( + e, + Effect::CreateSession { .. } | Effect::CreateWorktreeSession { .. } + ) + }) +} + +/// Conversation `LoadSession` must never inherit process-wide local stamp. +#[cfg(feature = "local-workspace")] +fn conversation_load_in_effects(effs: &[super::actions::Effect]) -> bool { + use super::actions::Effect; + effs.iter().any(|e| { + matches!( + e, + Effect::LoadSession { + chat_kind: true, + .. + } + ) + }) +} + +/// Apply history bypass (`chat_mode = false`) for load/restore/worktree-create. +#[cfg(feature = "local-workspace")] +pub(crate) fn welcome_history_build_bypass_applies( + effs: &[super::actions::Effect], + flag: bool, ) -> bool { - let flags = effects::SessionFlags { + use super::actions::Effect; + flag && effs.iter().any(|e| { + matches!( + e, + Effect::LoadSession { .. } + | Effect::RestoreAndLoadSession { .. } + | Effect::CreateWorktreeSession { .. } + ) + }) +} + +/// Whether this batch should clear the welcome history bypass flag. +#[cfg(feature = "local-workspace")] +pub(crate) fn welcome_history_build_bypass_consume( + effs: &[super::actions::Effect], + flag: bool, +) -> bool { + use super::actions::Effect; + flag && effs.iter().any(|e| { + matches!( + e, + Effect::LoadSession { .. } | Effect::CreateWorktreeSession { .. } + ) + }) +} + +/// Shared [`SessionFlags`] builder (interactive loop + leader-cluster). +pub(crate) fn session_flags_for_effects( + app: &mut AppView, + #[cfg_attr(not(feature = "local-workspace"), allow(unused_variables))] + effs: &[super::actions::Effect], +) -> effects::SessionFlags { + effects::SessionFlags { plan_mode: app.plan_mode, subagents: app.subagents, ask_user: app.ask_user, @@ -3030,13 +3732,54 @@ fn process_effects( app.default_yolo, matches!(app.current_ui.permission_mode.as_deref(), Some("auto")), ), - chat_mode: app.chat_mode, + chat_mode: { + #[cfg(feature = "local-workspace")] + { + if welcome_history_build_bypass_applies(effs, app.welcome_history_load_as_build) { + if welcome_history_build_bypass_consume(effs, app.welcome_history_load_as_build) + { + app.welcome_history_load_as_build = false; + } + false + } else { + app.chat_mode + } + } + #[cfg(not(feature = "local-workspace"))] + { + app.chat_mode + } + }, + #[cfg(feature = "local-workspace")] + local_workspace: { + if conversation_load_in_effects(effs) { + None // conversation resume is sandbox/gateway-owned + } else if welcome_oneshot_applies_to_effects(effs) { + match app.welcome_session_local_workspace.take() { + Some(one_shot) => one_shot, + None => crate::app::session_startup::active_local_workspace().unwrap_or(None), + } + } else { + crate::app::session_startup::active_local_workspace().unwrap_or(None) + } + }, screen_mode_label: Some(app.screen_mode.meta_label()), is_api_key_auth: app.is_api_key_auth, resume_local_miss: app.resume_local_miss.clone(), - }; + } +} + +/// Spawn effects into the task set. Returns `true` if the app should quit. +fn process_effects( + effs: Vec, + tasks: &mut JoinSet, + app: &mut AppView, + progress_tx: &tokio::sync::mpsc::UnboundedSender, +) -> bool { + let flags = session_flags_for_effects(app, &effs); for eff in effs { let (quit, meta) = effects::execute(eff, tasks, &app.acp_tx, &app.cwd, &flags, progress_tx); + // Install auth abort handle if the current auth state still matches. if let Some((seq, abort_handle)) = meta.auth_abort_handle && let super::app_view::AuthState::Authenticating { request_seq, @@ -3047,6 +3790,8 @@ fn process_effects( { *handle = Some(abort_handle); } + // Install URL-poll abort handle when the seq still matches (or is the + // current Authenticating attempt). Aborted in `abort_prior_auth`. if let Some((seq, abort_handle)) = meta.auth_url_poll_handle { let still_current = matches!( &app.auth_state, @@ -3063,10 +3808,85 @@ fn process_effects( } false } + #[cfg(test)] mod tests { use super::*; use crossterm::event::{KeyEvent, KeyEventState}; + + #[cfg(feature = "local-workspace")] + #[test] + fn welcome_oneshot_applies_to_create_worktree_session() { + use crate::app::actions::Effect; + use crate::app::agent::AgentId; + let worktree = Effect::CreateWorktreeSession { + agent_id: AgentId(0), + load_session_id: None, + label: None, + git_ref: None, + model_id: None, + preferred_session_id: None, + chat_kind: false, + }; + assert!(welcome_oneshot_applies_to_effects(std::slice::from_ref( + &worktree + ))); + assert!(!welcome_oneshot_applies_to_effects(&[])); + assert!(!welcome_oneshot_applies_to_effects(&[Effect::Quit])); + } + + #[cfg(feature = "local-workspace")] + #[test] + fn conversation_load_is_not_welcome_oneshot_or_local_stamp() { + use crate::app::actions::Effect; + use crate::app::agent::AgentId; + let load = Effect::LoadSession { + agent_id: AgentId(0), + session_id: "c1".into(), + session_cwd: None, + chat_kind: true, + }; + assert!(!welcome_oneshot_applies_to_effects(std::slice::from_ref( + &load + ))); + assert!(conversation_load_in_effects(std::slice::from_ref(&load))); + let build_load = Effect::LoadSession { + agent_id: AgentId(0), + session_id: "b1".into(), + session_cwd: None, + chat_kind: false, + }; + assert!(!conversation_load_in_effects(std::slice::from_ref( + &build_load + ))); + } + + #[cfg(feature = "local-workspace")] + #[test] + fn session_flags_consume_history_bypass_and_strip_conversation_stamp() { + use crate::app::actions::Effect; + use crate::app::agent::AgentId; + let mut app = crate::app::app_view::tests::test_app(); + app.chat_mode = true; + app.welcome_history_load_as_build = true; + let load = Effect::LoadSession { + agent_id: AgentId(0), + session_id: "c1".into(), + session_cwd: None, + chat_kind: true, + }; + let flags = session_flags_for_effects(&mut app, std::slice::from_ref(&load)); + assert!(!flags.chat_mode, "history bypass must clear chat_mode"); + assert!( + !app.welcome_history_load_as_build, + "LoadSession consumes the bypass" + ); + assert!( + flags.local_workspace.is_none(), + "conversation load must strip local stamp" + ); + } + #[test] fn tty_suspend_arm_stops_same_batch_before_later_ownership_changes() { let mut app = crate::app::app_view::tests::test_app(); @@ -3079,6 +3899,9 @@ mod tests { ); assert!(tty_suspend_armed(&app)); } + + // ── is_voice_chord ─────────────────────────────────────────────────── + #[test] fn is_voice_chord_press_exact_release_keycode() { use KeyEventKind::{Press, Release}; @@ -3096,15 +3919,23 @@ mod tests { KeyModifiers::CONTROL, KeyModifiers::NONE, ); + // Press: exact chord only — stray mods / bare Space don't fire (Thread 4). assert!(hit(sp, ctrl, Press) && hit(f8, none, Press)); assert!(!hit(sp, ctrl | KeyModifiers::ALT, Press)); assert!(!hit(f8, KeyModifiers::SHIFT, Press) && !hit(sp, none, Press)); + // Release: key alone — a bare Space release (Ctrl lifted first) matches so + // hold-to-talk can still stop (Thread 3); non-chord keys don't. assert!(hit(sp, none, Release) && hit(f8, none, Release)); assert!(!hit(KeyCode::Char('a'), none, Release)); } + + // ── voice_chord_action ─────────────────────────────────────────────── + #[test] fn voice_chord_action_cases() { use crate::app::actions::Action; + // (hold_mode, releases_reported, kind, listening, hold_owned) -> action + // tag, with the toggle-stop case being a past regression. let press = KeyEventKind::Press; let release = KeyEventKind::Release; let tag = |a: Option| match a { @@ -3115,10 +3946,15 @@ mod tests { _ => "other", }; let cases = [ + // hold + releases: press idle starts; release stops; press on a + // hold-owned session waits; press on a non-hold (/voice/toggle) + // session toggles off. ((true, true, press, false, false), "start"), ((true, true, release, true, true), "stop"), ((true, true, press, true, true), "none"), ((true, true, press, true, false), "toggle"), + // Non-hold (toggle mode or no reported releases): press toggles, + // release noops. ((false, false, press, false, false), "toggle"), ((false, false, release, true, false), "none"), ((true, false, release, true, false), "none"), @@ -3131,6 +3967,7 @@ mod tests { ); } } + /// Hold-owned events are claimed even with the setting off (a dropped /// release would wedge the mic open — past regression); otherwise presses /// honor the setting and bare releases are never claimed. @@ -3139,15 +3976,19 @@ mod tests { let press = KeyEventKind::Press; let repeat = KeyEventKind::Repeat; let release = KeyEventKind::Release; + // (kind, keybind_enabled, hold_owned) -> claimed let cases = [ + // Hold-owned: everything claimed, setting on or off. ((release, false, true), true), ((release, true, true), true), ((press, false, true), true), ((repeat, false, true), true), + // No hold: press/repeat follow the setting. ((press, true, false), true), ((press, false, false), false), ((repeat, true, false), true), ((repeat, false, false), false), + // No hold: a bare release is never ours (normal typing). ((release, true, false), false), ((release, false, false), false), ]; @@ -3159,11 +4000,15 @@ mod tests { ); } } + + // ── plan_reconnect_load ────────────────────────────────────────────── + #[test] fn plan_reconnect_load_requires_session_id() { let agent = crate::test_util::make_agent_view(None, "/work/project"); assert!(plan_reconnect_load(&agent, std::path::Path::new("/pager/cwd")).is_none()); } + /// The session's own cwd keys its on-disk storage — the pager cwd /// is only a fallback for agents without one. #[test] @@ -3172,10 +4017,12 @@ mod tests { let plan = plan_reconnect_load(&agent, std::path::Path::new("/pager/cwd")).unwrap(); assert_eq!(plan.session_id.0.as_ref(), "sess-1"); assert_eq!(plan.cwd, std::path::PathBuf::from("/work/worktree-a")); + let agent = crate::test_util::make_agent_view(Some("sess-1"), ""); let plan = plan_reconnect_load(&agent, std::path::Path::new("/pager/cwd")).unwrap(); assert_eq!(plan.cwd, std::path::PathBuf::from("/pager/cwd")); } + /// The reconnect cursor rides `_meta.cursor` when known; yolo mode /// always rides `_meta.yoloMode`. Auto rides `_meta.autoMode` per-agent. #[test] @@ -3187,20 +4034,28 @@ mod tests { plan.meta.get("cursor").is_none(), "no cursor key before any event was applied" ); + // autoMode is always set explicitly (false when not in auto) so the leader's + // capability injection can't re-enable Auto on reconnect. assert_eq!(plan.meta["autoMode"], serde_json::json!(false)); + agent.last_seen_event_id = Some("sess-1-42".into()); agent.session.yolo_mode = true; let plan = plan_reconnect_load(&agent, std::path::Path::new("/pager/cwd")).unwrap(); assert_eq!(plan.meta["yoloMode"], serde_json::json!(true)); assert_eq!(plan.meta["cursor"], serde_json::json!("sess-1-42")); } + #[test] fn plan_reconnect_load_meta_carries_auto_mode_from_session() { + // Auto rides `_meta.autoMode`, derived from THIS agent's own + // `auto_mode` (per-agent, symmetric with yolo) — not the global UI mirror. let mut agent = crate::test_util::make_agent_view(Some("sess-1"), "/work"); agent.session.auto_mode = true; let plan = plan_reconnect_load(&agent, std::path::Path::new("/pager/cwd")).unwrap(); assert_eq!(plan.meta["yoloMode"], serde_json::json!(false)); assert_eq!(plan.meta["autoMode"], serde_json::json!(true)); + + // Yolo wins: autoMode is explicitly false even if the session is in auto. let mut agent = crate::test_util::make_agent_view(Some("sess-1"), "/work"); agent.session.auto_mode = true; agent.session.yolo_mode = true; @@ -3208,6 +4063,7 @@ mod tests { assert_eq!(plan.meta["yoloMode"], serde_json::json!(true)); assert_eq!(plan.meta["autoMode"], serde_json::json!(false)); } + /// Multi-agent reconnect must seed each tab's `autoMode` from ITS OWN /// session, not a shared global mirror: an active Auto tab and a background /// Ask tab reconnect with `autoMode:true` and `autoMode:false` respectively. @@ -3216,9 +4072,12 @@ mod tests { let mut active = crate::test_util::make_agent_view(Some("sess-active"), "/work"); active.session.auto_mode = true; let background = crate::test_util::make_agent_view(Some("sess-bg"), "/work"); + // background.session.auto_mode stays false (Ask). + let active_plan = plan_reconnect_load(&active, std::path::Path::new("/pager/cwd")).unwrap(); let background_plan = plan_reconnect_load(&background, std::path::Path::new("/pager/cwd")).unwrap(); + assert_eq!(active_plan.meta["autoMode"], serde_json::json!(true)); assert_eq!( background_plan.meta["autoMode"], @@ -3226,11 +4085,13 @@ mod tests { "background Ask tab must reconnect with autoMode:false regardless of the active tab" ); } + #[test] fn reconnect_restores_dashboard_peek_before_replacing_scrollback() { use crate::scrollback::block::RenderBlock; use crate::views::dashboard::{DashboardRowId, DashboardState}; use indexmap::IndexMap; + let id = super::super::agent::AgentId(0); let mut agent = crate::test_util::make_agent_view(Some("sess-1"), "/work"); agent @@ -3248,11 +4109,16 @@ mod tests { .begin_peek_viewport(DashboardRowId::TopLevel(id), &mut agents); assert!(dashboard.as_ref().unwrap().peek_viewport.is_some()); assert!(agents[&id].scrollback.is_follow_mode()); + restore_dashboard_peek_before_reload(&mut dashboard, &mut agents); + assert!(dashboard.as_ref().unwrap().peek_viewport.is_none()); assert_eq!(agents[&id].scrollback.selected(), Some(0)); assert!(!agents[&id].scrollback.is_follow_mode()); } + + // ── reconnect_restore_outcome ──────────────────────────────────────── + /// The regression guard: one background tab fails, the active tab /// succeeds. The whole-reconnect flag goes false (toast says "failed"), /// but the active tab's OWN drain must still fire — a failed background tab @@ -3266,6 +4132,7 @@ mod tests { loads.insert(active, (true, None, None)); loads.insert(background, (false, None, None)); let pending = vec![active, background]; + let (all_restored, active_restored) = reconnect_restore_outcome(true, &pending, &loads, Some(active)); assert!( @@ -3277,6 +4144,7 @@ mod tests { "the active tab's own success still drains its queue" ); } + /// The active tab's OWN reload failed: its drain stays suppressed even /// though a background tab succeeded. #[test] @@ -3288,6 +4156,7 @@ mod tests { loads.insert(active, (false, None, None)); loads.insert(background, (true, None, None)); let pending = vec![active, background]; + let (all_restored, active_restored) = reconnect_restore_outcome(true, &pending, &loads, Some(active)); assert!(!all_restored); @@ -3296,6 +4165,7 @@ mod tests { "the active tab's own failure must block its drain" ); } + /// Single-agent behavior is preserved: the lone active tab succeeds → both /// flags true (toast "restored" + drain). #[test] @@ -3305,11 +4175,13 @@ mod tests { let mut loads = std::collections::HashMap::new(); loads.insert(active, (true, None, None)); let pending = vec![active]; + let (all_restored, active_restored) = reconnect_restore_outcome(true, &pending, &loads, Some(active)); assert!(all_restored); assert!(active_restored); } + /// A failed init (`init_ok == false`, empty `loads`) suppresses everything. #[test] fn reconnect_drain_blocked_when_init_failed() { @@ -3317,11 +4189,13 @@ mod tests { let active = AgentId(0); let loads = std::collections::HashMap::new(); let pending = vec![active]; + let (all_restored, active_restored) = reconnect_restore_outcome(false, &pending, &loads, Some(active)); assert!(!all_restored); assert!(!active_restored); } + /// No active agent (dashboard/welcome view): nothing to drain, even when /// every reloaded tab restored. #[test] @@ -3331,6 +4205,7 @@ mod tests { let mut loads = std::collections::HashMap::new(); loads.insert(background, (true, None, None)); let pending = vec![background]; + let (all_restored, active_restored) = reconnect_restore_outcome(true, &pending, &loads, None); assert!(all_restored); @@ -3339,9 +4214,11 @@ mod tests { "no active agent → no active-tab drain to fire" ); } + fn timed(event: Event, arrived_at: std::time::Instant) -> TimedInputEvent { TimedInputEvent { event, arrived_at } } + fn key_event(code: KeyCode, modifiers: KeyModifiers, kind: KeyEventKind) -> TimedInputEvent { TimedInputEvent::now(Event::Key(KeyEvent { code, @@ -3350,6 +4227,7 @@ mod tests { state: KeyEventState::NONE, })) } + fn scroll_event( kind: crossterm::event::MouseEventKind, arrived_at: std::time::Instant, @@ -3364,18 +4242,23 @@ mod tests { arrived_at, ) } + fn press(code: KeyCode) -> TimedInputEvent { key_event(code, KeyModifiers::NONE, KeyEventKind::Press) } + fn release(code: KeyCode) -> TimedInputEvent { key_event(code, KeyModifiers::NONE, KeyEventKind::Release) } + fn press_shift(code: KeyCode) -> TimedInputEvent { key_event(code, KeyModifiers::SHIFT, KeyEventKind::Press) } + fn press_ctrl(code: KeyCode) -> TimedInputEvent { key_event(code, KeyModifiers::CONTROL, KeyEventKind::Press) } + #[cfg(target_os = "linux")] fn mouse_event( kind: crossterm::event::MouseEventKind, @@ -3388,21 +4271,26 @@ mod tests { modifiers, })) } + #[test] fn park_input_reader_timeout_clears_stale_acknowledgement() { use std::sync::atomic::{AtomicBool, Ordering}; + let input_paused = AtomicBool::new(false); let reader_parked = AtomicBool::new(true); let acknowledged = park_input_reader(&input_paused, &reader_parked, Duration::ZERO); + assert!(!acknowledged); assert!(!reader_parked.load(Ordering::Acquire)); assert!(input_paused.load(Ordering::Acquire)); } + #[test] fn suspend_retry_gate_blocks_until_deadline() { let now = Instant::now(); let mut retry_after = None; let mut wait_reported = false; + assert!(defer_suspend_retry( &mut retry_after, &mut wait_reported, @@ -3412,6 +4300,8 @@ mod tests { assert_eq!(retry_after, Some(now + SUSPEND_RETRY_DELAY)); assert!(suspend_retry_ready(retry_after, now + SUSPEND_RETRY_DELAY)); assert!(wait_reported); + + // Mirrors the timer arm: expiry opens the gate for the next loop top. retry_after = None; assert!(suspend_retry_ready(retry_after, now)); assert!(!defer_suspend_retry( @@ -3422,17 +4312,22 @@ mod tests { assert_eq!(retry_after, Some(now + SUSPEND_RETRY_DELAY)); assert!(!suspend_retry_ready(retry_after, now)); } + #[test] fn suspend_timeout_requeues_request() { let mut pending = None; + requeue_after_suspend_timeout(&mut pending, "request"); + assert_eq!(pending, Some("request")); } + #[test] fn suspend_wait_feedback_is_reported_only_once_across_retries() { let now = Instant::now(); let mut retry_after = None; let mut reports = SuspendWaitReports::default(); + assert!(defer_suspend_retry( &mut retry_after, &mut reports.editor_reported, @@ -3444,6 +4339,7 @@ mod tests { &mut reports.editor_reported, now )); + reports.reset_missing(false, false); assert!(!reports.editor_reported); retry_after = None; @@ -3453,18 +4349,22 @@ mod tests { now )); } + #[test] fn editor_report_then_success_does_not_suppress_pager_first_timeout() { let now = Instant::now(); let mut retry_after = None; let mut reports = SuspendWaitReports::default(); + assert!(defer_suspend_retry( &mut retry_after, &mut reports.editor_reported, now )); + // The editor retry succeeds while the pager request remains pending. retry_after = None; reports.editor_reported = false; + assert!(defer_suspend_retry( &mut retry_after, &mut reports.pager_reported, @@ -3477,6 +4377,7 @@ mod tests { now )); } + #[test] fn suspend_wait_sink_is_mode_appropriate() { assert_eq!( @@ -3492,16 +4393,20 @@ mod tests { SuspendWaitSink::Toast ); } + #[test] fn suspend_wait_report_uses_system_block_in_minimal_mode() { use crate::scrollback::block::RenderBlock; + let mut app = crate::app::app_view::tests::test_app(); let id = crate::app::agent::AgentId(0); let agent = crate::test_util::make_agent_view(Some("session"), "/tmp"); app.agents.insert(id, agent); app.active_view = ActiveView::Agent(id); app.screen_mode = crate::app::ScreenMode::Minimal; + report_suspend_wait(&mut app, EDITOR_SUSPEND_WAIT); + let agent = app.agents.get(&id).expect("active agent"); let entry = agent.scrollback.last().expect("system block"); assert!(matches!( @@ -3510,6 +4415,7 @@ mod tests { )); assert!(agent.toast.is_none()); } + #[test] fn suspend_wait_report_uses_toast_outside_minimal_mode() { let mut app = crate::app::app_view::tests::test_app(); @@ -3518,7 +4424,9 @@ mod tests { app.agents.insert(id, agent); app.active_view = ActiveView::Agent(id); app.screen_mode = crate::app::ScreenMode::Inline; + report_suspend_wait(&mut app, EDITOR_SUSPEND_WAIT); + let agent = app.agents.get(&id).expect("active agent"); assert_eq!( agent.toast.as_ref().map(|(message, _)| message.as_str()), @@ -3526,18 +4434,22 @@ mod tests { ); assert!(agent.scrollback.last().is_none()); } + #[test] fn writer_failure_event_returns_original_error() { let error = writer_event_sequence(crate::render::draw::WriterEvent::Failed( std::io::Error::other("injected writer failure"), )) .expect_err("writer failure must terminate the event loop"); + assert_eq!(error.to_string(), "injected writer failure"); } + #[test] fn presenter_coalesces_until_ack() { let mut presenter = Presenter::new(); let mut draws = 0; + presenter.request(false); assert!(presenter.try_present(0, |_| draws += 1, || 1)); assert_eq!(presenter.in_flight_target, Some(1)); @@ -3547,22 +4459,27 @@ mod tests { } assert_eq!(draws, 1); assert!(presenter.dirty); + presenter.acknowledge(1); assert!(presenter.try_present(1, |_| draws += 1, || 2)); assert_eq!(draws, 2); assert_eq!(presenter.in_flight_target, Some(2)); } + #[test] fn presenter_no_output_does_not_wedge() { let mut presenter = Presenter::new(); presenter.request(false); + assert!(presenter.try_present(4, |_| {}, || 4)); assert_eq!(presenter.in_flight_target, None); assert!(!presenter.dirty); + presenter.request(false); assert!(presenter.try_present(4, |_| {}, || 5)); assert_eq!(presenter.in_flight_target, Some(5)); } + #[test] fn presenter_keeps_forced_repaint_sticky() { let mut presenter = Presenter { @@ -3572,11 +4489,13 @@ mod tests { presenter.request(false); presenter.request(true); let mut forced = false; + presenter.acknowledge(8); assert!(presenter.try_present(8, |force| forced = force, || 9)); assert!(forced); assert!(!presenter.force_full_repaint); } + #[test] fn presenter_immediate_ack_before_request_is_not_lost() { let mut presenter = Presenter { @@ -3585,30 +4504,37 @@ mod tests { }; presenter.acknowledge(3); presenter.request(false); + assert!(presenter.try_present(3, |_| {}, || 4)); assert_eq!(presenter.in_flight_target, Some(4)); } + #[test] fn presenter_later_ack_clears_target() { let mut presenter = Presenter { in_flight_target: Some(3), ..Presenter::new() }; + presenter.acknowledge(4); + assert_eq!(presenter.in_flight_target, None); } + #[test] fn presenter_waits_for_last_payload_in_turn() { let mut presenter = Presenter::new(); presenter.request(false); assert!(presenter.try_present(10, |_| {}, || 13)); presenter.request(false); + presenter.acknowledge(11); assert!(!presenter.try_present(13, |_| panic!("target not acknowledged"), || 14)); presenter.acknowledge(13); assert!(presenter.try_present(13, |_| {}, || 14)); assert_eq!(presenter.in_flight_target, Some(14)); } + #[test] fn timed_paste_uses_first_contributing_event() { let start = std::time::Instant::now(); @@ -3626,10 +4552,12 @@ mod tests { start + Duration::from_millis(8), ), ]; + let coalesced = coalesce_rapid_keys(events); assert_eq!(coalesced.len(), 1); assert_eq!(coalesced[0].arrived_at, start); assert_eq!(coalesced[0].event, Event::Paste("a\nb".to_owned())); + let fragments = vec![ timed(Event::Paste("a".to_owned()), start), timed( @@ -3645,9 +4573,11 @@ mod tests { assert_eq!(merged[0].arrived_at, start); assert_eq!(merged[0].event, Event::Paste("a\nb".to_owned())); } + #[test] fn delayed_scroll_batch_preserves_arrival_spacing_and_reversal() { use crossterm::event::MouseEventKind::{ScrollDown, ScrollUp}; + let mut app = crate::app::app_view::tests::test_app(); let start = std::time::Instant::now() + Duration::from_secs(1); app.scroll_state = Default::default(); @@ -3670,6 +4600,7 @@ mod tests { spaced.stream.expect("up stream active").avg_interval_ms, Some(8.0) ); + let routed = normalize_input_event(scroll_event(ScrollDown, start + Duration::from_millis(40))); let _ = app.handle_input_at_with_paste_provenance( @@ -3677,6 +4608,7 @@ mod tests { routed.arrived_at, routed.paste_provenance, ); + let snapshot = app .scroll_state .debug_snapshot(&app.scroll_config, start + Duration::from_millis(40)); @@ -3685,6 +4617,7 @@ mod tests { assert_eq!(stream.events, 1); assert_eq!(stream.gap_remaining_ms, 80); } + #[cfg(target_os = "linux")] #[test] fn unmodified_middle_down_reads_primary_once() { @@ -3695,18 +4628,21 @@ mod tests { x11_primary_available: true, ..Default::default() }); + let input = mouse_event( MouseEventKind::Down(MouseButton::Middle), KeyModifiers::NONE, ); let arrived_at = input.arrived_at; let normalized = normalize_input_event(input); + assert_eq!(normalized.event, Event::Paste("PRIMARY\nexact".to_owned())); assert_eq!(normalized.arrived_at, arrived_at); assert_eq!(normalized.paste_provenance, PasteProvenance::X11Primary); assert_eq!(crate::clipboard::primary_selection_read_call_count(), 1); crate::clipboard::clear_clipboard_probe_hook(); } + #[cfg(target_os = "linux")] #[test] fn nonqualifying_mouse_events_do_not_read_primary() { @@ -3716,6 +4652,7 @@ mod tests { x11_primary_available: true, ..Default::default() }); + let release = mouse_event(MouseEventKind::Up(MouseButton::Middle), KeyModifiers::NONE); let normalized = normalize_input_event(release.clone()); assert_eq!(normalized.event, release.event); @@ -3734,6 +4671,7 @@ mod tests { assert_eq!(crate::clipboard::primary_selection_read_call_count(), 0); crate::clipboard::clear_clipboard_probe_hook(); } + #[cfg(target_os = "linux")] #[test] fn empty_primary_preserves_original_middle_event() { @@ -3747,12 +4685,14 @@ mod tests { MouseEventKind::Down(MouseButton::Middle), KeyModifiers::NONE, ); + let normalized = normalize_input_event(middle.clone()); assert_eq!(normalized.event, middle.event); assert_eq!(normalized.paste_provenance, PasteProvenance::Terminal); assert_eq!(crate::clipboard::primary_selection_read_call_count(), 1); crate::clipboard::clear_clipboard_probe_hook(); } + #[test] fn coalesce_multiline_paste_without_bracketed_paste() { let events = vec![ @@ -3766,8 +4706,10 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste("ab\ncd".to_string())); } + #[test] fn coalesce_filters_release_events() { + // Press+Release pairs (Windows Terminal, Kitty) must not break runs. let events = vec![ press(KeyCode::Char('a')), release(KeyCode::Char('a')), @@ -3782,6 +4724,7 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste("ab\nc".to_string())); } + #[test] fn coalesce_preserves_shifted_chars() { let events = vec![ @@ -3796,6 +4739,7 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste("Hi\nBye".to_string())); } + #[test] fn coalesce_below_threshold_no_change() { let events = vec![press(KeyCode::Char('a')), press(KeyCode::Enter)]; @@ -3804,8 +4748,10 @@ mod tests { assert!(matches!(&result[0].event, Event::Key(ke) if ke.code == KeyCode::Char('a'))); assert!(matches!(&result[1].event, Event::Key(ke) if ke.code == KeyCode::Enter)); } + #[test] fn coalesce_no_enter_no_change() { + // No Enter in the run — no premature-send risk. let events = vec![ press(KeyCode::Char('h')), press(KeyCode::Char('e')), @@ -3819,8 +4765,10 @@ mod tests { assert!(matches!(&ev.event, Event::Key(_))); } } + #[test] fn coalesce_only_enters_no_change() { + // All-Enter runs must not coalesce (held Enter key repeat). let events = vec![ press(KeyCode::Enter), press(KeyCode::Enter), @@ -3830,6 +4778,7 @@ mod tests { let result = coalesce_rapid_keys(events); assert_eq!(result.len(), 4); } + #[test] fn coalesce_preserves_non_key_events() { let events = vec![ @@ -3845,6 +4794,7 @@ mod tests { assert_eq!(result[1].event, Event::Paste("a\nb".to_string())); assert!(matches!(&result[2].event, Event::Resize(100, 30))); } + #[test] fn coalesce_ctrl_key_breaks_run() { let events = vec![ @@ -3855,8 +4805,10 @@ mod tests { press(KeyCode::Char('d')), ]; let result = coalesce_rapid_keys(events); + // "ab" (2, no Enter) | Ctrl+C | "\nd" (2) — both runs below threshold. assert_eq!(result.len(), 5); } + #[test] fn coalesce_tabs_in_pasted_code() { let events = vec![ @@ -3870,6 +4822,7 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste("if\n\tx".to_string())); } + #[test] fn coalesce_exactly_at_threshold() { let events = vec![ @@ -3881,8 +4834,10 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste("a\nb".to_string())); } + #[test] fn coalesce_type_then_submit_not_coalesced() { + // Enter is the LAST event — "type + submit", not paste. let events = vec![ press(KeyCode::Char('a')), press(KeyCode::Char('b')), @@ -3893,8 +4848,10 @@ mod tests { assert_eq!(result.len(), 4); assert!(matches!(&result[3].event, Event::Key(ke) if ke.code == KeyCode::Enter)); } + #[test] fn fragmented_paste_merged_with_keys() { + // Event::Paste mixed with key events — merge into one paste. let events = vec![ TimedInputEvent::now(Event::Paste("real paste".into())), press(KeyCode::Char('a')), @@ -3905,6 +4862,7 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste("real pastea\nb".to_string())); } + #[test] fn coalesce_single_event_passthrough() { let events = vec![press(KeyCode::Enter)]; @@ -3912,13 +4870,18 @@ mod tests { assert_eq!(result.len(), 1); assert!(matches!(&result[0].event, Event::Key(_))); } + #[test] fn coalesce_empty_input() { let result = coalesce_rapid_keys(vec![]); assert!(result.is_empty()); } + + // ── Multi-newline coalescing tests ─────────────────────────────── + #[test] fn coalesce_three_lines() { + // "foo\nbar\nbaz" — 3 lines, 2 newlines. let events = vec![ press(KeyCode::Char('f')), press(KeyCode::Char('o')), @@ -3936,8 +4899,10 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste("foo\nbar\nbaz".to_string())); } + #[test] fn coalesce_four_lines_trailing_newline() { + // "a\nb\nc\nd\n" — 4 lines + trailing newline. let events = vec![ press(KeyCode::Char('a')), press(KeyCode::Enter), @@ -3952,16 +4917,21 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste("a\nb\nc\nd\n".to_string())); } + + // ── should_extend_for_paste tests ─────────────────────────────── + #[test] fn extend_triggered_with_single_pasteable_key() { let events = vec![press(KeyCode::Char('a'))]; assert!(should_extend_for_paste(&events)); } + #[test] fn extend_triggered_with_enter_key() { let events = vec![press(KeyCode::Enter)]; assert!(should_extend_for_paste(&events)); } + #[test] fn extend_not_triggered_with_bracketed_paste() { let events = vec![ @@ -3972,13 +4942,18 @@ mod tests { ]; assert!(!should_extend_for_paste(&events)); } + #[test] fn extend_not_triggered_with_only_non_pasteable() { let events = vec![TimedInputEvent::now(Event::Resize(80, 24))]; assert!(!should_extend_for_paste(&events)); } + + // ── merge_paste_fragments tests ───────────────────────────────── + #[test] fn merge_paste_and_key_fragments() { + // Fragmented bracketed paste: Event::Paste + loose key events. let events = vec![ TimedInputEvent::now(Event::Paste("hello\nwor".into())), press(KeyCode::Char('l')), @@ -3988,6 +4963,7 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste("hello\nworld".to_string())); } + #[test] fn merge_multiple_paste_fragments() { let events = vec![ @@ -3999,6 +4975,7 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste("aa\nbb\nc".to_string())); } + #[test] fn merge_preserves_non_key_events() { let events = vec![ @@ -4012,6 +4989,7 @@ mod tests { assert!(matches!(result[1].event, Event::Resize(80, 24))); assert_eq!(result[2].event, Event::Paste("x".to_string())); } + #[test] fn merge_skips_release_events() { let events = vec![ @@ -4023,6 +5001,7 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste("abc".to_string())); } + #[test] fn pure_paste_no_merge_needed() { let events = vec![TimedInputEvent::now(Event::Paste("hello\nworld".into()))]; @@ -4030,6 +5009,9 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste("hello\nworld".to_string())); } + + // ── is_pasteable_key_event filtering tests ───────────────────────── + #[test] fn pasteable_rejects_mouse_events() { use crossterm::event::{MouseButton, MouseEvent, MouseEventKind}; @@ -4048,20 +5030,24 @@ mod tests { }); assert!(!is_pasteable_key_event(&click)); } + #[test] fn pasteable_rejects_focus_events() { assert!(!is_pasteable_key_event(&Event::FocusGained)); assert!(!is_pasteable_key_event(&Event::FocusLost)); } + #[test] fn pasteable_rejects_release_events() { assert!(!is_pasteable_key_event(&release(KeyCode::Char('a')).event)); assert!(!is_pasteable_key_event(&release(KeyCode::Enter).event)); } + #[test] fn pasteable_rejects_resize() { assert!(!is_pasteable_key_event(&Event::Resize(80, 24))); } + #[test] fn pasteable_rejects_repeat_events() { let ev = Event::Key(KeyEvent { @@ -4072,6 +5058,7 @@ mod tests { }); assert!(!is_pasteable_key_event(&ev)); } + #[test] fn pasteable_accepts_valid_key_presses() { assert!(is_pasteable_key_event(&press(KeyCode::Char('a')).event)); @@ -4081,6 +5068,7 @@ mod tests { assert!(is_pasteable_key_event(&press(KeyCode::Enter).event)); assert!(is_pasteable_key_event(&press(KeyCode::Tab).event)); } + #[test] fn extend_not_triggered_with_only_mouse_and_focus() { use crossterm::event::{MouseEvent, MouseEventKind}; @@ -4095,6 +5083,7 @@ mod tests { ]; assert!(!should_extend_for_paste(&events)); } + #[test] fn extend_triggered_only_when_key_present_in_mixed_batch() { use crossterm::event::{MouseEvent, MouseEventKind}; @@ -4110,8 +5099,12 @@ mod tests { ]; assert!(should_extend_for_paste(&events)); } + #[test] fn coalesce_mouse_events_interleaved_with_paste_chars() { + // Simulates the batch produced by the fixed detect_paste: + // a key press followed by mouse events. The mouse events + // should not prevent the key from being processed. use crossterm::event::{MouseEvent, MouseEventKind}; let events = vec![ press(KeyCode::Char('a')), @@ -4129,13 +5122,17 @@ mod tests { })), ]; let result = coalesce_rapid_keys(events); + // Below coalesce threshold, all events pass through unchanged. assert_eq!(result.len(), 3); assert!(matches!(&result[0].event, Event::Key(ke) if ke.code == KeyCode::Char('a'))); assert!(matches!(&result[1].event, Event::Mouse(_))); assert!(matches!(&result[2].event, Event::Mouse(_))); } + #[test] fn coalesce_mouse_breaks_key_run_preserves_events() { + // A genuine paste batch that also collected mouse events. + // The paste chars should still coalesce; mouse events are preserved. use crossterm::event::{MouseEvent, MouseEventKind}; let events = vec![ press(KeyCode::Char('a')), @@ -4150,12 +5147,22 @@ mod tests { press(KeyCode::Char('c')), ]; let result = coalesce_rapid_keys(events); + // The mouse event breaks the key run: [a, b, Enter] (3 keys, but + // Enter is last in that sub-run → no char after Enter → not coalesced), + // then [mouse], then [c] (1 key). assert_eq!(result.len(), 5); } + + // ── Windows path-shape coalescing (drag-drop without bracketed paste) ─ + // + // Windows-gated: the path-shape branch only exists on Windows + // (other platforms reliably get bracketed paste for drag-drop). + #[cfg(target_os = "windows")] fn press_run(text: &str) -> Vec { text.chars().map(|c| press(KeyCode::Char(c))).collect() } + /// Smoke test across every anchor variant the branch should match: /// drive-letter (both separators), UNC, Unix absolute, `file://`, /// and the Windows-Terminal-quoted form for paths with spaces. @@ -4175,24 +5182,26 @@ mod tests { assert_eq!(result[0].event, Event::Paste(input.to_string())); } } + /// Below-threshold path-shape (< 8 chars) and non-path prose of any /// length must NOT coalesce — keep typed editing intact. #[cfg(target_os = "windows")] #[test] fn coalesce_path_shape_rejects_short_or_non_path() { - let short = "/foo.tx"; + let short = "/foo.tx"; // 7 chars, below PATH_COALESCE_THRESHOLD assert!( coalesce_rapid_keys(press_run(short)) .iter() .all(|e| matches!(e.event, Event::Key(_))) ); - let prose = "helloworld"; + let prose = "helloworld"; // 10 chars, no path anchor assert!( coalesce_rapid_keys(press_run(prose)) .iter() .all(|e| matches!(e.event, Event::Key(_))) ); } + /// `:` in a US-layout drive-letter path arrives as Shift+`;`; /// `is_pasteable_key_event` accepts SHIFT so the run must assemble /// cleanly. @@ -4206,6 +5215,9 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0].event, Event::Paste(r"C:\foo.png".to_string())); } + + // ── make_run_result exit info ──────────────────────────────────────── + /// App focused on an agent (session `test-session`) with a seeded /// prompt → prompt → response exchange in its scrollback. fn seeded_quit_app(screen_mode: crate::app::ScreenMode) -> AppView { @@ -4221,6 +5233,7 @@ mod tests { scrollback.push_block(RenderBlock::agent_message("Pinned the seed.\nSecond line.")); app } + #[test] fn make_run_result_fullscreen_quit_builds_summary() { let app = seeded_quit_app(crate::app::ScreenMode::Fullscreen); @@ -4228,6 +5241,7 @@ mod tests { assert_eq!(info.session_id, "test-session"); assert!(!info.minimal); let summary = info.summary.expect("summary on fullscreen quit"); + // Deliberate: title comes from the first prompt, last_prompt from the newest. assert_eq!(summary.title, "fix the flaky CI test"); assert_eq!( summary.last_prompt.as_deref(), @@ -4235,6 +5249,7 @@ mod tests { ); assert_eq!(summary.last_response.as_deref(), Some("Pinned the seed.")); } + #[test] fn make_run_result_unanswered_prompt_omits_stale_response() { use crate::scrollback::block::RenderBlock; @@ -4253,19 +5268,23 @@ mod tests { summary.last_prompt.as_deref(), Some("now rerun the whole suite") ); + // The earlier reply answered an older prompt — it must not appear here. assert!(summary.last_response.is_none()); } + #[test] fn make_run_result_inline_and_minimal_quits_omit_summary() { let app = seeded_quit_app(crate::app::ScreenMode::Inline); let info = make_run_result(&app).exit_info.expect("agent exit info"); assert!(info.summary.is_none()); assert!(!info.minimal); + let app = seeded_quit_app(crate::app::ScreenMode::Minimal); let info = make_run_result(&app).exit_info.expect("agent exit info"); assert!(info.summary.is_none()); assert!(info.minimal); } + #[test] fn make_run_result_empty_session_omits_summary() { let mut app = crate::app::app_view::tests::test_app_with_agent(); @@ -4273,6 +5292,7 @@ mod tests { let info = make_run_result(&app).exit_info.expect("agent exit info"); assert!(info.summary.is_none()); } + #[test] fn make_run_result_non_agent_views_have_no_exit_info() { for view in [ActiveView::Welcome, ActiveView::AgentDashboard] { diff --git a/crates/codegen/xai-grok-pager/src/app/leader_cluster/mod.rs b/crates/codegen/xai-grok-pager/src/app/leader_cluster/mod.rs index f49e9a2..58c7ae0 100644 --- a/crates/codegen/xai-grok-pager/src/app/leader_cluster/mod.rs +++ b/crates/codegen/xai-grok-pager/src/app/leader_cluster/mod.rs @@ -29,19 +29,13 @@ //! seams), where env is set before any process-global's first touch. //! //! Unix-only: the leader transport here is a unix socket. -use super::actions::{Action, TaskResult}; -use super::agent::AgentState; -use super::agent_view::AgentView; -use super::app_view::{AppView, AuthState, TrustState}; -use super::{acp_handler, dispatch, effects}; -use crate::acp::leader_bridge::bridge_channels; -use crate::acp::model_state::ModelState; -use crate::scrollback::block::RenderBlock; -use agent_client_protocol as acp; + use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicUsize}; use std::time::Duration; + +use agent_client_protocol as acp; use tempfile::TempDir; use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; @@ -52,8 +46,19 @@ use xai_grok_shell::leader::{ LeaderServerControlState, LeaderServerMetadata, ReconnectPolicy, run_leader_server, }; use xai_grok_test_support::MockInferenceServer; + +use super::actions::{Action, TaskResult}; +use super::agent::AgentState; +use super::agent_view::AgentView; +use super::app_view::{AppView, AuthState, TrustState}; +use super::{acp_handler, dispatch, effects}; +use crate::acp::leader_bridge::bridge_channels; +use crate::acp::model_state::ModelState; +use crate::scrollback::block::RenderBlock; + const PUMP_TICK: Duration = Duration::from_millis(10); const TURN_BUDGET: Duration = Duration::from_secs(60); + /// Await a bring-up step with a hard budget so an on-demand run that hangs /// names its phase instead of parking until the test-runner kill. async fn bounded(what: &str, fut: impl std::future::Future) -> T { @@ -61,11 +66,13 @@ async fn bounded(what: &str, fut: impl std::future::Future) -> T .await .unwrap_or_else(|_| panic!("leader-cluster bring-up timed out: {what}")) } + /// The grok home the agent actually persisted under: `grok_home()` is /// process-cached, so an earlier test in this binary may have pinned it. fn effective_grok_home() -> PathBuf { xai_grok_config::grok_home() } + /// Concatenated agent-message text across a view's scrollback (copy of the /// acp_handler tests' helper; that one is test-mod private). fn agent_message_text(view: &AgentView) -> String { @@ -79,6 +86,7 @@ fn agent_message_text(view: &AgentView) -> String { } out } + /// One pager client: a full `AppView` behind the production leader bridge. struct ClusterClient { app: AppView, @@ -91,6 +99,7 @@ struct ClusterClient { /// generation bumps after a leader kill/respawn. status_rx: Option>, } + impl ClusterClient { /// Drain everything currently ready (inbound ACP + finished tasks). /// Returns whether anything was processed. @@ -110,31 +119,18 @@ impl ClusterClient { } progressed } + fn drain_pending_effects(&mut self) { if !self.app.pending_effects.is_empty() { let effs = std::mem::take(&mut self.app.pending_effects); self.process_effects(effs); } } + /// The event loop's `process_effects`, minus terminal/auth-handle wiring /// (that fn is event_loop-private; this mirrors its body). fn process_effects(&mut self, effs: Vec) { - let flags = effects::SessionFlags { - plan_mode: self.app.plan_mode, - subagents: self.app.subagents, - ask_user: self.app.ask_user, - restore_code: self.app.restore_code, - agent_override: self.app.agent_override.clone(), - yolo_mode: self.app.default_yolo, - auto_mode: dispatch::effective_auto( - self.app.default_yolo, - matches!(self.app.current_ui.permission_mode.as_deref(), Some("auto")), - ), - chat_mode: self.app.chat_mode, - screen_mode_label: Some(self.app.screen_mode.meta_label()), - is_api_key_auth: self.app.is_api_key_auth, - resume_local_miss: self.app.resume_local_miss.clone(), - }; + let flags = super::event_loop::session_flags_for_effects(&mut self.app, &effs); for eff in effs { let (_quit, _meta) = effects::execute( eff, @@ -147,17 +143,20 @@ impl ClusterClient { } self.drain_pending_effects(); } + /// Dispatch a user action and run its effects. fn act(&mut self, action: Action) { let effs = dispatch::dispatch(action, &mut self.app); self.process_effects(effs); } + /// Pump until `pred(app)` holds, within [`TURN_BUDGET`]. No fixed sleeps /// beyond the pump tick; panics with `what` on expiry. Single-client sugar /// over [`pump_clients_until`] so there is exactly one pump loop. async fn pump_until(&mut self, what: &str, pred: impl Fn(&AppView) -> bool) { pump_clients_until(&mut [self], what, |clients| pred(&clients[0].app)).await; } + /// The most recently created agent view (scenarios add tabs in order). fn latest_agent(&self) -> &AgentView { self.app @@ -166,6 +165,7 @@ impl ClusterClient { .last() .expect("client has no agent view yet") } + fn agent_for_session(&self, sid: &str) -> &AgentView { self.app .agents @@ -178,6 +178,7 @@ impl ClusterClient { }) .unwrap_or_else(|| panic!("no agent view for session {sid}")) } + /// Create a new session through the real dispatch → effect → agent path. async fn new_session(&mut self) -> String { self.act(Action::NewSession); @@ -195,6 +196,7 @@ impl ClusterClient { .0 .to_string() } + /// Attach to an existing session (viewer path) and wait for the replay to /// land. async fn load_session(&mut self, sid: &str) { @@ -211,6 +213,7 @@ impl ClusterClient { }) .await; } + /// Drive one full turn on the active agent and wait until it lands /// (sentinel visible + agent back to Idle). async fn run_turn(&mut self, prompt: &str, sentinel: &str) { @@ -224,10 +227,12 @@ impl ClusterClient { }) .await; } + fn sever(self) { self.bridge_cancel.cancel(); } } + /// Pump several clients until `pred` holds across them, within /// [`TURN_BUDGET`]; panics with `what` on expiry. async fn pump_clients_until( @@ -250,6 +255,7 @@ async fn pump_clients_until( tokio::time::sleep(PUMP_TICK).await; } } + /// The cluster: leader server + real agent, plus knobs to kill/respawn the /// leader generation under the same socket path. struct PagerLeaderCluster { @@ -275,15 +281,18 @@ struct PagerLeaderCluster { _env: Vec, _grok_home: TempDir, } + impl PagerLeaderCluster { /// Stand up the cluster. Callers MUST be `#[serial_test::serial(GROK_HOME)]` /// (env mutation) and run inside a current-thread `LocalSet`. async fn start() -> Self { let _ = rustls::crypto::ring::default_provider().install_default(); + let server = MockInferenceServer::start().await.expect("mock server"); let grok_home = TempDir::new().unwrap(); let workdir = TempDir::new().unwrap(); let sock_path = grok_home.path().join("leader-cluster.sock"); + let env = vec![ crate::test_util::EnvVarGuard::set("GROK_HOME", grok_home.path()), crate::test_util::EnvVarGuard::set("GROK_CLI_CHAT_PROXY_BASE_URL", server.url()), @@ -296,12 +305,15 @@ impl PagerLeaderCluster { // connect_or_spawn) to this cluster's socket. crate::test_util::EnvVarGuard::set(LEADER_SOCKET_ENV, &sock_path), ]; + + // Hold the flock for the cluster's lifetime (see field doc). let mut flock = LeaderLock::new(""); assert!( flock.try_acquire().expect("acquire cluster flock"), "cluster flock unexpectedly held" ); flock.write_pid().expect("stamp cluster flock"); + let client_count = Arc::new(AtomicUsize::new(0)); let mut cluster = Self { sock_path, @@ -318,6 +330,7 @@ impl PagerLeaderCluster { cluster.spawn_leader_generation().await; cluster } + /// Bind a fresh leader-server generation at the fixed socket path and /// wire a fresh REAL agent behind it. async fn spawn_leader_generation(&mut self) { @@ -326,11 +339,15 @@ impl PagerLeaderCluster { let (response_tx, response_rx) = tokio::sync::mpsc::unbounded_channel::(); let cancel = CancellationToken::new(); self.server_cancel = cancel.clone(); + let control_state = LeaderServerControlState::new(LeaderServerMetadata { pid: std::process::id(), socket_path: self.sock_path.clone(), lock_path: self.sock_path.with_extension("lock"), ws_url_suffix: String::new(), + // MUST be the client-side comparison source (xai_grok_version), not + // this crate's version: a reconnecting client evicts strictly-older + // leaders, and "evict" here would signal THIS test process. leader_binary_version: xai_grok_version::VERSION.to_string(), }); let sock_for_server = self.sock_path.clone(); @@ -355,17 +372,20 @@ impl PagerLeaderCluster { ) .await; })); + generation_tasks.extend(xai_grok_shell::leader::in_process::spawn_agent( acp_rx, response_tx, )); self.generation_tasks = generation_tasks; + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); while !self.sock_path.exists() && tokio::time::Instant::now() < deadline { tokio::time::sleep(Duration::from_millis(20)).await; } assert!(self.sock_path.exists(), "leader socket never bound"); } + /// Kill the current leader generation (server + agent die together, like /// a real leader process crash) and wait for the socket to vanish. async fn kill_leader(&mut self) { @@ -374,19 +394,32 @@ impl PagerLeaderCluster { while self.sock_path.exists() && tokio::time::Instant::now() < deadline { tokio::time::sleep(Duration::from_millis(20)).await; } + // Fail HERE if the old generation never released the socket: its late + // shutdown cleanup would otherwise delete the respawned generation's + // fresh socket from under it (same-path race), which surfaces as a + // confusing reconnect-budget expiry downstream. assert!( !self.sock_path.exists(), "old leader generation never released the socket" ); + // Abort + drain the generation's agent/bridge tasks (the server task + // has already run its socket cleanup above). Channel-closure teardown + // is only eventual; without this drain an old agent task could still + // be running against the same GROK_HOME when the next generation's + // agent starts — two writers on one updates.jsonl, the corruption + // class the real leader's flock prevents. for task in self.generation_tasks.drain(..) { task.abort(); let _ = task.await; } + // The next generation's agent must re-authenticate its ACP surface. self.authenticated = false; } + async fn respawn_leader(&mut self) { self.spawn_leader_generation().await; } + /// Connect a pager client. With `reconnect: true` the bridge gets a real /// `LeaderReconnector` (socket pinned via `GROK_LEADER_SOCKET`, flock held /// by the cluster, so reconnects always adopt the in-process server). @@ -406,6 +439,7 @@ impl PagerLeaderCluster { .await .expect("cluster client connect"); let (leader_tx, leader_rx) = conn.into_channels(); + let cancel = CancellationToken::new(); let (reconnector, status_rx) = if reconnect { let (status_tx, status_rx) = LeaderReconnector::status_channel(); @@ -426,6 +460,7 @@ impl PagerLeaderCluster { } else { (None, None) }; + let bridge = bridge_channels( leader_tx, leader_rx, @@ -436,6 +471,8 @@ impl PagerLeaderCluster { .expect("bridge spawn"); let tx = bridge.channel.tx; let rx = bridge.channel.rx; + + // Same handshake the pager performs after bridging (spawn path). let _init: acp::InitializeResponse = bounded( "initialize", acp_send( @@ -476,12 +513,14 @@ impl PagerLeaderCluster { .expect("authenticate through bridge"); self.authenticated = true; } + let mut app = AppView::new(tx, ModelState::default(), Vec::new()); app.leader_mode = true; app.auth_state = AuthState::Done; app.trust_state = TrustState::Done; app.project_picker_shown = true; app.cwd = self.workdir.path().to_path_buf(); + let (progress_tx, progress_rx) = tokio::sync::mpsc::unbounded_channel(); ClusterClient { app, @@ -493,6 +532,7 @@ impl PagerLeaderCluster { status_rx, } } + /// Inference request count (chat/responses/messages only), for /// no-turn-was-re-driven invariants. fn inference_request_count(&self) -> usize { @@ -507,15 +547,20 @@ impl PagerLeaderCluster { .count() } } + impl Drop for PagerLeaderCluster { fn drop(&mut self) { self.server_cancel.cancel(); + // Best-effort (Drop cannot await): stop the generation's tasks so they + // never outlive the env guards / temp dirs dropping right after. for task in self.generation_tasks.drain(..) { task.abort(); } } } + fn occurrences(haystack: &str, needle: &str) -> usize { haystack.matches(needle).count() } + mod scenarios; diff --git a/crates/codegen/xai-grok-pager/src/app/mod.rs b/crates/codegen/xai-grok-pager/src/app/mod.rs index ffbbd2a..6405bbd 100644 --- a/crates/codegen/xai-grok-pager/src/app/mod.rs +++ b/crates/codegen/xai-grok-pager/src/app/mod.rs @@ -652,6 +652,19 @@ pub async fn run( { anyhow::bail!("{err}"); } + #[cfg(feature = "local-workspace")] + { + let lw = session_startup::resolve_local_workspace_config( + args.chat(), + args.local_workspace(), + args.local_workspace_attach(), + args.local_workspace_cwd(), + )?; + if let Some(ref cfg) = lw { + session_startup::emit_local_workspace_startup_ux(cfg)?; + } + session_startup::set_active_local_workspace(lw)?; + } let intent = args .session_startup_intent() .map_err(|e| anyhow::anyhow!("{e}"))?; @@ -2018,6 +2031,47 @@ mod tests { fn cli_chat_flag_rejected_without_feature() { assert!(try_parse_pager(&["grok-pager", "--chat"]).is_err()); } + #[cfg(feature = "local-workspace")] + #[test] + fn cli_local_workspace_attach_requires_chat() { + assert!( + try_parse_pager(&["grok-pager", "--local-workspace-attach=srv"]).is_err(), + "attach without --chat must clap-error" + ); + let args = + try_parse_pager(&["grok-pager", "--chat", "--local-workspace-attach=srv"]).unwrap(); + assert_eq!(args.local_workspace_attach(), Some("srv")); + } + #[cfg(feature = "local-workspace")] + #[test] + fn cli_local_workspace_own_conflicts_with_attach() { + assert!( + try_parse_pager(&[ + "grok-pager", + "--chat", + "--local-workspace=/tmp/a", + "--local-workspace-attach=srv", + ]) + .is_err(), + "own + attach must clap-conflict" + ); + } + #[cfg(feature = "local-workspace")] + #[test] + fn cli_local_workspace_cwd_requires_chat() { + assert!(try_parse_pager(&["grok-pager", "--local-workspace-cwd=/tmp/a"]).is_err()); + let args = try_parse_pager(&[ + "grok-pager", + "--chat", + "--local-workspace-attach=srv", + "--local-workspace-cwd=/tmp/repo", + ]) + .unwrap(); + assert_eq!( + args.local_workspace_cwd(), + Some(std::path::Path::new("/tmp/repo")) + ); + } #[test] fn cli_local_workspace_flags_rejected_without_feature() { assert!(try_parse_pager(&["grok-pager", "--local-workspace-attach=srv"]).is_err()); diff --git a/crates/codegen/xai-grok-pager/src/app/modals.rs b/crates/codegen/xai-grok-pager/src/app/modals.rs index eb1d590..f6feb5d 100644 --- a/crates/codegen/xai-grok-pager/src/app/modals.rs +++ b/crates/codegen/xai-grok-pager/src/app/modals.rs @@ -1022,41 +1022,19 @@ impl AgentView { vim_normal_first: crate::appearance::cache::load_vim_mode(), }; - // Delete-confirmation flow: `d` arms a confirmation on the - // focused row, then `y` confirms and `n` (or any other key) - // cancels. y/n are intercepted here — before the picker - // handler — only while armed, so the rest of the time `y` - // keeps its normal meaning (copy the session id). - if pending_delete.is_some() - && let crossterm::event::Event::Key(k) = ev - && k.kind == KeyEventKind::Press - && k.modifiers.is_empty() - { - match k.code { - crossterm::event::KeyCode::Char('y') => { - // Confirm. The cwd was captured when the row was - // armed, so this can't be foiled by an async - // picker-list update (e.g. a deep-search result) - // landing between `d` and `y`. - if let Some((source, session_id, cwd)) = pending_delete.take() { - return InputOutcome::Action(Action::DeleteSession { - source, - session_id, - cwd, - }); - } - return InputOutcome::Changed; - } - crossterm::event::KeyCode::Char('n') => { - *pending_delete = None; - return InputOutcome::Changed; - } - _ => { - // Any other key cancels, then falls through to its - // normal handling below. - *pending_delete = None; - } + match crate::views::session_picker::handle_pending_delete_key(pending_delete, ev) { + crate::views::session_picker::PendingDeleteKey::Confirm(pd) => { + return InputOutcome::Action(Action::DeleteSession { + source: pd.source, + session_id: pd.session_id, + cwd: pd.cwd, + }); } + crate::views::session_picker::PendingDeleteKey::Cancel => { + return InputOutcome::Changed; + } + crate::views::session_picker::PendingDeleteKey::Disarmed + | crate::views::session_picker::PendingDeleteKey::NotArmed => {} } if let crossterm::event::Event::Key(key) = ev @@ -1082,7 +1060,12 @@ impl AgentView { }); } - match handle_picker_input(ev, state, entry_count, &config) { + let selected_before = state.selected; + let outcome = handle_picker_input(ev, state, entry_count, &config); + if pending_delete.is_some() && state.selected != selected_before { + *pending_delete = None; + } + match outcome { PickerOutcome::Selected(i) => { match entry_map.get(i).and_then(|e| e.as_ref()) { Some(PickerItem::Fuzzy { original_index }) => { @@ -1225,28 +1208,13 @@ impl AgentView { InputOutcome::Action(Action::CycleSessionSourceFilter) } PickerOutcome::Action('d') => { - // Arm a delete confirmation on the highlighted row, - // capturing source, id, and cwd now (the row is present at - // this moment) so the `y` confirm can't be foiled by an - // async picker-list update. `y` confirms / `n` cancels - // on the next key press (intercepted above). *pending_delete = - match entry_map.get(state.selected).and_then(|e| e.as_ref()) { - Some(PickerItem::Fuzzy { original_index }) => entries - .as_ref() - .and_then(|e| e.get(*original_index)) - .filter(|entry| { - !crate::app::foreign_sessions::is_foreign_picker_source( - &entry.source, - ) - }) - .map(|e| (e.source.clone(), e.id.clone(), e.cwd.clone())), - Some(PickerItem::Content { hit_index }) => content_results - .as_ref() - .and_then(|h| h.get(*hit_index)) - .map(|h| ("local".into(), h.session_id.clone(), h.cwd.clone())), - None => None, - }; + crate::views::session_picker::pending_delete_from_selection( + state.selected, + &entry_map, + entries.as_deref(), + content_results.as_deref(), + ); InputOutcome::Changed } PickerOutcome::NonSelectableClick(_) @@ -2403,7 +2371,7 @@ mod session_picker_delete_tests { use crate::app::agent_view::test_fixtures::make_agent; use crate::app::app_view::{InputOutcome, SessionPickerEntry}; use crate::views::modal::ActiveModal; - use crossterm::event::{Event, KeyCode, KeyEvent, KeyModifiers}; + use crossterm::event::{Event, KeyCode, KeyEvent, KeyModifiers, MouseEvent, MouseEventKind}; fn entry(id: &str) -> SessionPickerEntry { SessionPickerEntry { @@ -2447,9 +2415,9 @@ mod session_picker_delete_tests { fn pending(agent: &AgentView) -> Option { match agent.active_modal.as_ref() { - Some(ActiveModal::SessionPicker { pending_delete, .. }) => pending_delete - .as_ref() - .map(|(_, session_id, _)| session_id.clone()), + Some(ActiveModal::SessionPicker { pending_delete, .. }) => { + pending_delete.as_ref().map(|pd| pd.session_id.clone()) + } _ => None, } } @@ -2509,6 +2477,20 @@ mod session_picker_delete_tests { ); } + #[test] + fn mouse_move_keeps_pending_delete() { + let mut agent = make_agent(); + open_picker(&mut agent, vec![entry("s0"), entry("s1")]); + agent.handle_palette_or_arg_input(&key('d')); + agent.handle_palette_or_arg_input(&Event::Mouse(MouseEvent { + kind: MouseEventKind::Moved, + column: 0, + row: 0, + modifiers: KeyModifiers::NONE, + })); + assert_eq!(pending(&agent).as_deref(), Some("s0")); + } + #[test] fn y_without_armed_confirmation_does_not_delete() { let mut agent = make_agent(); diff --git a/crates/codegen/xai-grok-pager/src/app/session_startup.rs b/crates/codegen/xai-grok-pager/src/app/session_startup.rs index 11347a2..b7bf010 100644 --- a/crates/codegen/xai-grok-pager/src/app/session_startup.rs +++ b/crates/codegen/xai-grok-pager/src/app/session_startup.rs @@ -44,6 +44,9 @@ pub struct DeferredStartupActions { pub prompt: Option, pub open_dashboard: bool, pub pending_chat: bool, + /// Welcome history local-disk bypass persisted across the startup gate. + #[cfg(feature = "local-workspace")] + pub history_load_as_build: bool, } impl DeferredStartupActions { pub fn is_empty(&self) -> bool { @@ -300,9 +303,321 @@ pub fn chat_mode_flag_conflict( } None } +/// Env: enable local workspace without CLI flags (`1`). Mode defaults to `own` +/// unless `GROK_CHAT_LOCAL_WORKSPACE_MODE` / attach server id is set. +#[cfg(feature = "local-workspace")] +pub const GROK_CHAT_LOCAL_WORKSPACE_ENV: &str = "GROK_CHAT_LOCAL_WORKSPACE"; +#[cfg(feature = "local-workspace")] +pub const GROK_CHAT_LOCAL_WORKSPACE_CWD_ENV: &str = "GROK_CHAT_LOCAL_WORKSPACE_CWD"; +#[cfg(feature = "local-workspace")] +pub const GROK_CHAT_LOCAL_WORKSPACE_MODE_ENV: &str = "GROK_CHAT_LOCAL_WORKSPACE_MODE"; +#[cfg(feature = "local-workspace")] +pub const GROK_CHAT_LOCAL_WORKSPACE_SERVER_ID_ENV: &str = "GROK_CHAT_LOCAL_WORKSPACE_SERVER_ID"; +#[cfg(feature = "local-workspace")] +pub const GROK_CHAT_LOCAL_WORKSPACE_ALLOW_HOME_ENV: &str = "GROK_CHAT_LOCAL_WORKSPACE_ALLOW_HOME"; +/// Skip interactive first-run confirm (still prints the banner). +#[cfg(feature = "local-workspace")] +pub const GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV: &str = "GROK_CHAT_LOCAL_WORKSPACE_ACK"; +/// Startup banner / first-run copy. +#[cfg(feature = "local-workspace")] +pub const LOCAL_WORKSPACE_BANNER: &str = + "Local workspace runs tools on this machine (FS confined to )."; +#[cfg(feature = "local-workspace")] +pub const LOCAL_WORKSPACE_ATTACH_NEEDS_SERVER_ID: &str = "local-workspace attach requires --local-workspace-attach= \ + (or GROK_CHAT_LOCAL_WORKSPACE_SERVER_ID)"; +#[cfg(feature = "local-workspace")] +pub const LOCAL_WORKSPACE_REQUIRES_CHAT: &str = "local-workspace flags/env require --chat"; +#[cfg(feature = "local-workspace")] +pub const LOCAL_WORKSPACE_HOME_DENIED: &str = + "local-workspace cwd may not be / or $HOME unless GROK_CHAT_LOCAL_WORKSPACE_ALLOW_HOME=1"; +#[cfg(feature = "local-workspace")] +pub const LOCAL_WORKSPACE_HITL_HINT: &str = "Permission prompts for local workspace tools apply to your machine. \ + Local workspace replaces the chat sandbox."; +#[cfg(feature = "local-workspace")] +pub const LOCAL_WORKSPACE_ACK_REQUIRED: &str = + "local-workspace requires interactive confirm, GROK_CHAT_LOCAL_WORKSPACE_ACK=1, or an ack file"; +/// Declared advertised tool ids for attach FS-only check (comma-separated). +#[cfg(feature = "local-workspace")] +pub const GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS_ENV: &str = + "GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS"; +#[cfg(feature = "local-workspace")] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LocalWorkspaceMode { + Own, + Attach, +} +#[cfg(feature = "local-workspace")] +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LocalWorkspaceConfig { + pub mode: LocalWorkspaceMode, + pub cwd: Option, + pub server_id: Option, +} +#[cfg(feature = "local-workspace")] +static ACTIVE_LOCAL_WORKSPACE: std::sync::Mutex> = + std::sync::Mutex::new(None); +#[cfg(feature = "local-workspace")] +pub fn set_active_local_workspace(cfg: Option) -> anyhow::Result<()> { + let mut guard = ACTIVE_LOCAL_WORKSPACE.lock().map_err(|_| { + anyhow::anyhow!("local-workspace intent mutex poisoned; refuse attach (fail closed)") + })?; + tracing::info!( + target: crate::views::welcome::workspace_mode::WORKSPACE_MODE_LOG, + event = if cfg.is_some() { + "process_stamp_set" + } else { + "process_stamp_cleared" + }, + mode = cfg.as_ref().map(|c| format!("{:?}", c.mode)), + server_id = cfg.as_ref().and_then(|c| c.server_id.as_deref()), + cwd = cfg.as_ref().and_then(|c| c.cwd.as_ref().map(|p| p.display().to_string())), + "local-workspace process-wide intent stamp" + ); + *guard = cfg; + Ok(()) +} +#[cfg(feature = "local-workspace")] +pub fn active_local_workspace() -> anyhow::Result> { + ACTIVE_LOCAL_WORKSPACE + .lock() + .map(|g| g.clone()) + .map_err(|_| { + anyhow::anyhow!("local-workspace intent mutex poisoned; refuse attach (fail closed)") + }) +} +#[cfg(not(feature = "local-workspace"))] pub fn active_local_workspace() -> anyhow::Result> { Ok(None) } +#[cfg(feature = "local-workspace")] +fn env_truthy(name: &str) -> bool { + std::env::var(name) + .ok() + .is_some_and(|v| matches!(v.trim(), "1" | "true" | "TRUE" | "yes" | "YES")) +} +#[cfg(feature = "local-workspace")] +fn env_nonempty(name: &str) -> Option { + std::env::var(name) + .ok() + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) +} +/// Resolve CLI > env local-workspace intent (own or attach). +/// +/// Returns `Ok(None)` when local workspace is not requested. +#[cfg(feature = "local-workspace")] +pub fn resolve_local_workspace_config( + chat: bool, + cli_own: Option>, + cli_attach: Option<&str>, + cli_cwd: Option<&std::path::Path>, +) -> anyhow::Result> { + let env_enable = env_truthy(GROK_CHAT_LOCAL_WORKSPACE_ENV); + let env_mode = env_nonempty(GROK_CHAT_LOCAL_WORKSPACE_MODE_ENV); + let env_server_id = env_nonempty(GROK_CHAT_LOCAL_WORKSPACE_SERVER_ID_ENV); + let env_cwd = env_nonempty(GROK_CHAT_LOCAL_WORKSPACE_CWD_ENV).map(std::path::PathBuf::from); + let cli_attach = cli_attach.map(str::trim).filter(|s| !s.is_empty()); + let cli_requested = cli_own.is_some() || cli_attach.is_some(); + let env_requested = env_enable || env_mode.is_some() || env_server_id.is_some(); + if !cli_requested && !env_requested { + return Ok(None); + } + if !chat { + anyhow::bail!("{LOCAL_WORKSPACE_REQUIRES_CHAT}"); + } + let mode = if cli_attach.is_some() { + LocalWorkspaceMode::Attach + } else if cli_own.is_some() { + LocalWorkspaceMode::Own + } else if let Some(ref m) = env_mode { + match m.as_str() { + "attach" => LocalWorkspaceMode::Attach, + "own" => LocalWorkspaceMode::Own, + other => { + anyhow::bail!( + "invalid {GROK_CHAT_LOCAL_WORKSPACE_MODE_ENV}={other:?}; expected own|attach" + ) + } + } + } else if env_server_id.is_some() { + LocalWorkspaceMode::Attach + } else { + LocalWorkspaceMode::Own + }; + let cwd = cli_cwd + .map(std::path::Path::to_path_buf) + .or_else(|| cli_own.and_then(|inner| inner.map(std::path::Path::to_path_buf))) + .or(env_cwd) + .unwrap_or_else(|| { + std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")) + }); + let cwd = if cwd.is_absolute() { + cwd + } else { + std::env::current_dir() + .unwrap_or_else(|_| std::path::PathBuf::from(".")) + .join(cwd) + }; + let cwd = validate_local_workspace_cwd(&cwd)?; + match mode { + LocalWorkspaceMode::Own => Ok(Some(LocalWorkspaceConfig { + mode, + cwd: Some(cwd), + server_id: None, + })), + LocalWorkspaceMode::Attach => { + let server_id = cli_attach + .map(str::to_owned) + .or(env_server_id) + .filter(|s| !s.is_empty()); + let Some(server_id) = server_id else { + anyhow::bail!("{LOCAL_WORKSPACE_ATTACH_NEEDS_SERVER_ID}"); + }; + ensure_attach_fs_only_toolset(&server_id)?; + Ok(Some(LocalWorkspaceConfig { + mode, + cwd: Some(cwd), + server_id: Some(server_id), + })) + } + } +} +/// Canonicalize `path` and enforce the `/` + `$HOME` denylist. +/// +/// Returns the canonical directory so callers stamp/persist what was actually +/// checked (symlinks / `..` must not diverge from validation). +#[cfg(feature = "local-workspace")] +pub fn validate_local_workspace_cwd(path: &std::path::Path) -> anyhow::Result { + let abs = if path.is_absolute() { + path.to_path_buf() + } else { + std::env::current_dir() + .unwrap_or_else(|_| std::path::PathBuf::from(".")) + .join(path) + }; + let canon = abs.canonicalize().map_err(|e| { + anyhow::anyhow!( + "local workspace cwd must exist and be canonicalizable: {}: {e}", + abs.display() + ) + })?; + if !canon.is_dir() { + anyhow::bail!( + "local workspace cwd must be an existing directory: {}", + canon.display() + ); + } + if env_truthy(GROK_CHAT_LOCAL_WORKSPACE_ALLOW_HOME_ENV) { + return Ok(canon); + } + if canon == std::path::Path::new("/") { + anyhow::bail!("{LOCAL_WORKSPACE_HOME_DENIED}"); + } + if let Some(home_path) = dirs::home_dir().or_else(|| std::env::var_os("HOME").map(Into::into)) { + let home_canon = home_path.canonicalize().unwrap_or(home_path); + if canon == home_canon { + anyhow::bail!("{LOCAL_WORKSPACE_HOME_DENIED}"); + } + } + Ok(canon) +} +/// Banner + first-run confirm for local-workspace own/attach. +/// +/// Skip confirm only with `GROK_CHAT_LOCAL_WORKSPACE_ACK=1` or a prior ack file. +/// Non-TTY without ACK refuses (fail closed). +#[cfg(feature = "local-workspace")] +pub fn emit_local_workspace_startup_ux(cfg: &LocalWorkspaceConfig) -> anyhow::Result<()> { + use std::io::IsTerminal; + emit_local_workspace_startup_ux_with(cfg, std::io::stdin().is_terminal()) +} +/// Testable UX gate: `stdin_is_terminal` is injected. +#[cfg(feature = "local-workspace")] +pub fn emit_local_workspace_startup_ux_with( + cfg: &LocalWorkspaceConfig, + stdin_is_terminal: bool, +) -> anyhow::Result<()> { + let cwd_display = cfg + .cwd + .as_ref() + .map(|p| p.display().to_string()) + .unwrap_or_else(|| "".to_string()); + let banner = LOCAL_WORKSPACE_BANNER.replace("", &cwd_display); + eprintln!("{banner}"); + eprintln!("{LOCAL_WORKSPACE_HITL_HINT}"); + if local_workspace_ack_satisfied() { + return Ok(()); + } + if !stdin_is_terminal { + anyhow::bail!("{LOCAL_WORKSPACE_ACK_REQUIRED}"); + } + eprint!("Continue with local workspace on this machine? [y/N] "); + use std::io::Write; + let _ = std::io::stderr().flush(); + let mut line = String::new(); + std::io::stdin().read_line(&mut line)?; + let ok = matches!(line.trim(), "y" | "Y" | "yes" | "YES"); + if !ok { + anyhow::bail!("local workspace cancelled"); + } + write_local_workspace_ack(); + Ok(()) +} +/// True when ACK env or ack file already authorizes local workspace. +#[cfg(feature = "local-workspace")] +pub fn local_workspace_ack_satisfied() -> bool { + if env_truthy(GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV) { + return true; + } + local_workspace_ack_path().is_some_and(|p| p.is_file()) +} +/// Persist the first-run local-workspace ACK file (best-effort). +#[cfg(feature = "local-workspace")] +pub fn write_local_workspace_ack() { + if let Some(path) = local_workspace_ack_path() { + if let Some(parent) = path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let _ = std::fs::write(path, "1\n"); + } +} +/// Fail closed unless advertised tools are FS-only. +/// +/// Until diag exposes a real tool catalog, attach trusts operator attestation +/// via `GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS` (comma-separated ids). +/// Unset / empty → refuse. +#[cfg(feature = "local-workspace")] +pub fn ensure_attach_fs_only_toolset(_server_id: &str) -> anyhow::Result<()> { + let advertised = probe_advertised_tool_ids(); + let refs: Option> = advertised + .as_ref() + .map(|ids| ids.iter().map(String::as_str).collect()); + crate::app::effects::reject_non_fs_only_advertised_tools(refs.as_deref()) + .map_err(|e| anyhow::anyhow!("{e}")) +} +/// Operator-attested advertised tool ids for attach (env only; no fake diag probe). +#[cfg(feature = "local-workspace")] +pub fn probe_advertised_tool_ids() -> Option> { + let raw = env_nonempty(GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS_ENV)?; + let ids: Vec = raw + .split(',') + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect(); + Some(ids) +} +#[cfg(feature = "local-workspace")] +fn local_workspace_ack_path() -> Option { + let home = std::env::var("GROK_HOME") + .ok() + .filter(|s| !s.trim().is_empty()) + .map(std::path::PathBuf::from) + .or_else(|| { + dirs::home_dir() + .or_else(|| std::env::var_os("HOME").map(Into::into)) + .map(|h| h.join(".grok")) + })?; + Some(home.join("local_workspace_ack")) +} /// Conservative shape check for a chat-mode `--resume ` passthrough. /// /// The id skips disk/GCS resolution and flows to the gateway, but it is also @@ -1333,4 +1648,200 @@ mod tests { } } } + #[cfg(feature = "local-workspace")] + fn advertised_tools_env() -> xai_grok_test_support::EnvGuard { + xai_grok_test_support::EnvGuard::set( + GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS_ENV, + "workspace.fs_list,workspace.fs_read_file,workspace.fs_write_file,workspace.fs_exists,workspace.fs_delete_file,workspace.put_files,workspace.get_files", + ) + } + #[cfg(feature = "local-workspace")] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS)] + #[test] + fn resolve_local_workspace_attach_from_cli() { + let _env = advertised_tools_env(); + let tmp = tempfile::tempdir().unwrap(); + let cfg = resolve_local_workspace_config(true, None, Some("srv-dogfood"), Some(tmp.path())) + .unwrap() + .expect("attach config"); + assert_eq!(cfg.mode, LocalWorkspaceMode::Attach); + assert_eq!(cfg.server_id.as_deref(), Some("srv-dogfood")); + let canon = tmp.path().canonicalize().unwrap(); + assert_eq!(cfg.cwd.as_deref(), Some(canon.as_path())); + } + #[cfg(feature = "local-workspace")] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS)] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_SERVER_ID)] + #[test] + fn resolve_local_workspace_empty_cli_attach_falls_back_to_env() { + let _env = advertised_tools_env(); + let _sid = xai_grok_test_support::EnvGuard::set( + GROK_CHAT_LOCAL_WORKSPACE_SERVER_ID_ENV, + "srv-from-env", + ); + let tmp = tempfile::tempdir().unwrap(); + let cfg = resolve_local_workspace_config(true, None, Some(""), Some(tmp.path())) + .unwrap() + .expect("empty CLI attach should fall back to env server id"); + assert_eq!(cfg.mode, LocalWorkspaceMode::Attach); + assert_eq!(cfg.server_id.as_deref(), Some("srv-from-env")); + } + #[cfg(feature = "local-workspace")] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_CWD)] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE)] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_MODE)] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_SERVER_ID)] + #[test] + fn resolve_local_workspace_cwd_only_is_not_a_request() { + let tmp = tempfile::tempdir().unwrap(); + let _cwd = xai_grok_test_support::EnvGuard::set( + GROK_CHAT_LOCAL_WORKSPACE_CWD_ENV, + tmp.path().to_str().unwrap(), + ); + let _enable = xai_grok_test_support::EnvGuard::unset(GROK_CHAT_LOCAL_WORKSPACE_ENV); + let _mode = xai_grok_test_support::EnvGuard::unset(GROK_CHAT_LOCAL_WORKSPACE_MODE_ENV); + let _sid = xai_grok_test_support::EnvGuard::unset(GROK_CHAT_LOCAL_WORKSPACE_SERVER_ID_ENV); + let cfg = resolve_local_workspace_config(true, None, None, Some(tmp.path())).unwrap(); + assert!( + cfg.is_none(), + "cwd-only CLI/env must not activate local workspace: {cfg:?}" + ); + } + #[cfg(feature = "local-workspace")] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS)] + #[test] + fn resolve_local_workspace_own_from_cli() { + let _env = advertised_tools_env(); + let tmp = tempfile::tempdir().unwrap(); + let cfg = resolve_local_workspace_config(true, Some(Some(tmp.path())), None, None) + .unwrap() + .expect("own config"); + assert_eq!(cfg.mode, LocalWorkspaceMode::Own); + assert!( + cfg.server_id.is_none(), + "own leaves server_id to supervisor" + ); + let canon = tmp.path().canonicalize().unwrap(); + assert_eq!(cfg.cwd.as_deref(), Some(canon.as_path())); + } + #[cfg(feature = "local-workspace")] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS)] + #[test] + fn resolve_local_workspace_own_env_defaults() { + let _env = advertised_tools_env(); + let _enable = xai_grok_test_support::EnvGuard::set(GROK_CHAT_LOCAL_WORKSPACE_ENV, "1"); + let _mode = xai_grok_test_support::EnvGuard::unset(GROK_CHAT_LOCAL_WORKSPACE_MODE_ENV); + let _sid = xai_grok_test_support::EnvGuard::unset(GROK_CHAT_LOCAL_WORKSPACE_SERVER_ID_ENV); + let cwd = tempfile::tempdir().unwrap(); + let _cwd = xai_grok_test_support::EnvGuard::set( + GROK_CHAT_LOCAL_WORKSPACE_CWD_ENV, + cwd.path().to_str().unwrap(), + ); + let cfg = resolve_local_workspace_config(true, None, None, None) + .unwrap() + .expect("env own"); + assert_eq!(cfg.mode, LocalWorkspaceMode::Own); + assert!(cfg.server_id.is_none()); + let canon = cwd.path().canonicalize().unwrap(); + assert_eq!(cfg.cwd.as_deref(), Some(canon.as_path())); + } + #[cfg(feature = "local-workspace")] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS)] + #[test] + fn resolve_local_workspace_requires_chat() { + let _env = advertised_tools_env(); + let err = resolve_local_workspace_config(false, None, Some("srv"), None).unwrap_err(); + assert!( + err.to_string().contains("require --chat"), + "unexpected: {err}" + ); + } + #[cfg(feature = "local-workspace")] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS)] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ALLOW_HOME)] + #[serial_test::serial(HOME)] + #[serial_test::serial(USERPROFILE)] + #[test] + fn resolve_local_workspace_defaults_cwd_and_denies_home() { + let _tools = xai_grok_test_support::EnvGuard::set( + GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS_ENV, + "workspace.fs_list", + ); + let _allow = + xai_grok_test_support::EnvGuard::unset(GROK_CHAT_LOCAL_WORKSPACE_ALLOW_HOME_ENV); + let home = tempfile::tempdir().unwrap(); + let home_str = home.path().to_str().unwrap(); + let _home = xai_grok_test_support::EnvGuard::set("HOME", home_str); + let _userprofile = xai_grok_test_support::EnvGuard::set("USERPROFILE", home_str); + let err = + resolve_local_workspace_config(true, None, Some("srv"), Some(home.path())).unwrap_err(); + assert!(err.to_string().contains("ALLOW_HOME"), "unexpected: {err}"); + } + #[cfg(feature = "local-workspace")] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS)] + #[test] + fn resolve_local_workspace_refuses_uncheckable_toolset() { + let _tools = + xai_grok_test_support::EnvGuard::unset(GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS_ENV); + let tmp = tempfile::tempdir().unwrap(); + let err = + resolve_local_workspace_config(true, None, Some("srv"), Some(tmp.path())).unwrap_err(); + assert!( + err.to_string().contains("uncheckable") || err.to_string().contains("FS-only"), + "unexpected: {err}" + ); + } + #[cfg(feature = "local-workspace")] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS)] + #[test] + fn resolve_local_workspace_refuses_non_fs_toolset() { + let _tools = xai_grok_test_support::EnvGuard::set( + GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS_ENV, + "workspace.fs_list,workspace.bash", + ); + let tmp = tempfile::tempdir().unwrap(); + let err = + resolve_local_workspace_config(true, None, Some("srv"), Some(tmp.path())).unwrap_err(); + assert!(err.to_string().contains("FS-only"), "unexpected: {err}"); + assert!( + err.to_string().contains("workspace.bash"), + "unexpected: {err}" + ); + } + #[cfg(feature = "local-workspace")] + #[test] + fn local_workspace_banner_mentions_local_machine() { + assert!(LOCAL_WORKSPACE_BANNER.contains("on this machine")); + assert!(LOCAL_WORKSPACE_HITL_HINT.contains("your machine")); + assert!(LOCAL_WORKSPACE_HITL_HINT.contains("replaces the chat sandbox")); + } + #[cfg(feature = "local-workspace")] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ACK)] + #[serial_test::serial(GROK_HOME)] + #[test] + fn local_workspace_non_tty_requires_ack() { + let _ack = xai_grok_test_support::EnvGuard::unset(GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV); + let home = tempfile::tempdir().unwrap(); + let _home = + xai_grok_test_support::EnvGuard::set("GROK_HOME", home.path().to_str().unwrap()); + let cfg = LocalWorkspaceConfig { + mode: LocalWorkspaceMode::Attach, + cwd: Some(std::path::PathBuf::from("/tmp/repo")), + server_id: Some("srv".into()), + }; + let err = emit_local_workspace_startup_ux_with(&cfg, false).unwrap_err(); + assert!( + err.to_string().contains("ACK") || err.to_string().contains("ack"), + "unexpected: {err}" + ); + } + #[cfg(feature = "local-workspace")] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ALLOW_HOME)] + #[test] + fn validate_local_workspace_cwd_denies_root() { + let _allow = + xai_grok_test_support::EnvGuard::unset(GROK_CHAT_LOCAL_WORKSPACE_ALLOW_HOME_ENV); + let err = validate_local_workspace_cwd(std::path::Path::new("/")).unwrap_err(); + assert!(err.to_string().contains("ALLOW_HOME"), "{err}"); + } } diff --git a/crates/codegen/xai-grok-pager/src/app/xt_filter.rs b/crates/codegen/xai-grok-pager/src/app/xt_filter.rs index fd35ecd..1d65604 100644 --- a/crates/codegen/xai-grok-pager/src/app/xt_filter.rs +++ b/crates/codegen/xai-grok-pager/src/app/xt_filter.rs @@ -12,7 +12,7 @@ use super::event_loop::{TimedInputEvent, is_bare_esc_press}; const XT_ARM_WINDOW: Duration = Duration::from_secs(5); /// How long a held partial reply waits for its remaining fragments before -/// being resolved (pi-mono uses 150ms). +/// being resolved (other terminal UI stacks use 150ms). pub(super) const XT_FRAGMENT_TIMEOUT: Duration = Duration::from_millis(150); /// Total bound on one hold, so a terminal trickling valid payload chars diff --git a/crates/codegen/xai-grok-pager/src/pty_wrap.rs b/crates/codegen/xai-grok-pager/src/pty_wrap.rs index 3f10d08..0ae701f 100644 --- a/crates/codegen/xai-grok-pager/src/pty_wrap.rs +++ b/crates/codegen/xai-grok-pager/src/pty_wrap.rs @@ -66,7 +66,8 @@ pub(crate) fn run_wrapped_command(program: &str, args: &[String]) -> Result cmd.env("GROK_OSC52_SINK", "1"); cmd.env("LC_GROK_OSC52_SINK", "1"); - // Spawn child in the PTY slave. + // Not session-scoped: this is the wrapped process itself. + #[allow(clippy::disallowed_methods)] let mut child = pair.slave.spawn_command(cmd)?; // Drop the slave so we get EOF when child exits. drop(pair.slave); diff --git a/crates/codegen/xai-grok-pager/src/views/dashboard/mod.rs b/crates/codegen/xai-grok-pager/src/views/dashboard/mod.rs index c9e31b3..fade0cf 100644 --- a/crates/codegen/xai-grok-pager/src/views/dashboard/mod.rs +++ b/crates/codegen/xai-grok-pager/src/views/dashboard/mod.rs @@ -38,7 +38,7 @@ pub use render::{ }; pub use row::{ DashboardRow, RowBadge, build_rows, build_rows_with_roster, classify_subagent, - classify_top_level, sort_rows, + classify_top_level, roster_activity_to_state, sort_rows, }; pub use state::{ DashboardDispatchMode, DashboardRowId, DashboardState, Filter, FilterValue, Focusable, diff --git a/crates/codegen/xai-grok-pager/src/views/dashboard/render.rs b/crates/codegen/xai-grok-pager/src/views/dashboard/render.rs index 9b60e27..2e76dab 100644 --- a/crates/codegen/xai-grok-pager/src/views/dashboard/render.rs +++ b/crates/codegen/xai-grok-pager/src/views/dashboard/render.rs @@ -153,6 +153,13 @@ pub fn render_dashboard( home, roster, ); + // Chat-conversation roster rows can't be deleted from the dashboard + // yet — record them so the `[✗]` and Ctrl+X arm both skip them. + state.conversation_row_ids = roster + .iter() + .filter(|e| e.origin.kind == "conversation") + .map(|e| e.session_id.clone()) + .collect(); state.reanchor_selection(&rows); // DO NOT GC pinned/reorder at render time. The old @@ -592,6 +599,7 @@ fn render_dashboard_banner( use ratatui::widgets::{Block, Borders, Widget}; state.row_rects.clear(); + state.row_delete_rects.clear(); state.section_rects.clear(); if area.area() == 0 || area.height < 3 { return; @@ -1548,6 +1556,7 @@ fn render_rows( state: &mut DashboardState, ) { state.row_rects.clear(); + state.row_delete_rects.clear(); state.section_rects.clear(); state.idle_overflow_rect = None; if area.area() == 0 { @@ -2157,7 +2166,7 @@ fn render_row( rect: Rect, theme: &Theme, row: &DashboardRow, - state: &DashboardState, + state: &mut DashboardState, ) { if rect.area() == 0 { return; @@ -2293,25 +2302,57 @@ fn render_row( Style::default().bg(bg).fg(icon_color), ); - // Age column — reserve up to 8 cells on the right edge (to fit - // "just now"). Uses coarse buckets: just now / m / h / d / mo / y. + let armed_delete = state.armed_delete_row_ref(); + let show_delete = !row.is_more_placeholder + && !row.id.is_subagent() + && row.state.allows_delete() + && !state.row_is_conversation(&row.id) + && (state.hovered_row.as_ref() == Some(&row.id) || armed_delete == Some(&row.id)); + let delete_label = crate::glyphs::ballot_x_button(); + let delete_w = UnicodeWidthStr::width(delete_label) as u16; let age = format_time_ago(row.last_change_at.elapsed().unwrap_or_default()); let age_str = format!("{age:>6}"); let age_w = UnicodeWidthStr::width(age_str.as_str()) as u16; - let age_x = rect.x + rect.width.saturating_sub(age_w + 1); - if age_x > content_start_x { - buf.set_string( - age_x, - title_y, - &age_str, - Style::default().bg(bg).fg(theme.gray), - ); + let right_w = if show_delete { delete_w } else { age_w }; + let right_x = rect.x + rect.width.saturating_sub(right_w + 1); + if right_x > content_start_x { + if show_delete { + let fg = if state.hovered_delete.as_ref() == Some(&row.id) + || armed_delete == Some(&row.id) + { + theme.accent_error + } else { + theme.text_secondary + }; + buf.set_string( + right_x, + title_y, + delete_label, + Style::default().bg(bg).fg(fg), + ); + state.row_delete_rects.push(( + row.id.clone(), + Rect { + x: right_x, + y: title_y, + width: delete_w, + height: 1, + }, + )); + } else { + buf.set_string( + right_x, + title_y, + &age_str, + Style::default().bg(bg).fg(theme.gray), + ); + } } // Title text: `{label}` (bright) + ` · {subtitle}` (dim) + // optional `[badge]` chips for failed / pinned. // Trimmed to fit between the icon and the age column. - let title_avail = age_x.saturating_sub(content_start_x).saturating_sub(2); + let title_avail = right_x.saturating_sub(content_start_x).saturating_sub(2); let mut cx = content_start_x; if title_avail > 0 { let label_style = Style::default().bg(bg).fg(if row.is_more_placeholder { @@ -2353,9 +2394,9 @@ fn render_row( // Subtitle: ` · xai my-branch-2 worktree`. if let Some(sub) = row.subtitle.as_deref() - && cx + 4 < age_x + && cx + 4 < right_x { - let remaining = age_x.saturating_sub(cx).saturating_sub(2) as usize; + let remaining = right_x.saturating_sub(cx).saturating_sub(2) as usize; let sub_str = format!(" \u{00B7} {sub}"); let sub_trunc = truncate_str(&sub_str, remaining); let sub_w = UnicodeWidthStr::width(&sub_trunc[..]) as u16; @@ -2387,7 +2428,7 @@ fn render_row( }; let chip = format!(" [{label}]"); let cw = UnicodeWidthStr::width(chip.as_str()) as u16; - if cx + cw + 1 < age_x { + if cx + cw + 1 < right_x { buf.set_string( cx, title_y, @@ -2458,6 +2499,7 @@ fn render_narrow_rows( state: &mut DashboardState, ) { state.row_rects.clear(); + state.row_delete_rects.clear(); state.section_rects.clear(); state.idle_overflow_rect = None; if area.area() == 0 { @@ -2619,7 +2661,18 @@ fn render_narrow_rows( let indent_w = UnicodeWidthStr::width(indent.as_str()) as u16; let gap_after_marker = 1u16; let chrome = marker_w + gap_after_marker + indent_w + icon_w + 1; - let label = truncate_str(&row.label, body_width.saturating_sub(chrome) as usize); + let armed_here = state.armed_delete_row_ref() == Some(&row.id); + let show_delete = !row.is_more_placeholder + && !row.id.is_subagent() + && row.state.allows_delete() + && !state.row_is_conversation(&row.id) + && (hovered || armed_here); + let delete_label = crate::glyphs::ballot_x_button(); + let delete_w = UnicodeWidthStr::width(delete_label) as u16; + let label_budget = body_width + .saturating_sub(chrome) + .saturating_sub(if show_delete { delete_w + 1 } else { 0 }); + let label = truncate_str(&row.label, label_budget as usize); let line = format!("{marker} {indent}{icon} {label}"); buf.set_string( area.x, @@ -2627,6 +2680,18 @@ fn render_narrow_rows( line, Style::default().fg(theme.text_primary).bg(bg), ); + if show_delete && body_width > chrome + delete_w { + let dx = area.x + body_width.saturating_sub(delete_w); + let fg = if state.hovered_delete.as_ref() == Some(&row.id) || armed_here { + theme.accent_error + } else { + theme.text_secondary + }; + buf.set_string(dx, y, delete_label, Style::default().fg(fg).bg(bg)); + state + .row_delete_rects + .push((row.id.clone(), Rect::new(dx, y, delete_w, 1))); + } } if !row.is_more_placeholder { state.row_rects.push((row.id.clone(), line_rect)); @@ -3303,19 +3368,11 @@ fn render_file_search_dropdown_for( /// (no inline approve/reject yet — punted per the user's note /// "maybe its just easier to hit enter and go details view"; /// the dashboard is intentionally a navigator, not a permission UI). -/// - Anything else → `Enter:open · Ctrl+x:stop|close · ?:shortcuts`. +/// - Anything else → `Enter:open · Ctrl+x:stop|delete · ?:shortcuts`. /// /// The Ctrl+x chip label follows the selected agent's state: `stop` /// for an agent with a live turn (Working, or NeedsInput — paused but -/// still running, so the first Ctrl+x cancels), `close` for an idle / -/// quiet one. -/// -/// The ↑/↓ nav chip is intentionally omitted from every state — the -/// list is obviously arrow-navigable, and dropping it frees space so -/// the Ctrl+x chip stays visible while an agent is selected. -/// -/// Stop-confirm still routes through `with_pending` so the canonical -/// `press again to close this session` message takes over. +/// still running, so the first Ctrl+x cancels), `delete` otherwise. #[allow(clippy::too_many_arguments)] fn render_footer( buf: &mut Buffer, @@ -3358,27 +3415,31 @@ fn render_footer( return; } - // Only paint the "press again" hint while the confirm window is - // actually live — the dispatcher re-arms (rather than closes) on a - // press after [`super::state::STOP_CONFIRM_WINDOW`], so an expired - // confirm must not keep claiming the footer (e.g. after a mouse - // click moved the selection without a keypress to disarm it). - let stop_confirm_live = state - .stop_confirm - .as_ref() - .is_some_and(|(_, t)| t.elapsed() < super::state::STOP_CONFIRM_WINDOW); - if stop_confirm_live { - let stop_key = registry - .find(crate::actions::ActionId::DashboardStop) - .map(|d| d.default_key) - .unwrap_or_else(|| key!('x', CONTROL)); - let pending = PendingHint { - shortcut: stop_key, - label: "close this session", - }; - ShortcutsBar::new(&[]) - .with_pending(Some(pending)) - .render(inner, buf); + // A live delete-confirm owns the footer: `y`/`n` when the list is + // focused, else the second-`Ctrl+X` "press again" hint. An expired arm + // falls through to the normal hints. + if state.armed_delete_row_ref().is_some() { + if state.list_focused { + let hints = vec![ + HintItem::new(key!('y'), "confirm delete"), + HintItem::new(key!('n'), "cancel"), + ]; + ShortcutsBar::new(&hints) + .compact(4, None) + .render(inner, buf); + } else { + let stop_key = registry + .find(crate::actions::ActionId::DashboardStop) + .map(|d| d.default_key) + .unwrap_or_else(|| key!('x', CONTROL)); + let pending = PendingHint { + shortcut: stop_key, + label: "delete this session", + }; + ShortcutsBar::new(&[]) + .with_pending(Some(pending)) + .render(inner, buf); + } return; } @@ -3410,23 +3471,16 @@ fn render_footer( return; } - // A selected Inactive row is roster-only (owned by another pager - // process, never loaded here) — there's nothing running to stop, - // so every branch below suppresses its `stop` chip. - let stoppable = selected_state != Some(RowState::Inactive); - - // Ctrl+x cancels the live turn for a busy agent, else closes the - // session — mirroring `dispatch_dashboard_stop` (cancel-if-running, - // else close). A `NeedsInput` row keeps a paused-but-running turn (the - // permission/Q&A prompt suspends it, never idles it), so its first - // Ctrl+x cancels too — label it `stop`, not `close`. + let show_ctrl_x = selected_state.is_some_and(|s| { + matches!(s, RowState::Working | RowState::NeedsInput) || s.allows_delete() + }); let stop_label = if matches!( selected_state, Some(RowState::Working | RowState::NeedsInput) ) { "stop" } else { - "close" + "delete" }; // Overview list focused (via Tab) — navigation hints: arrows / j-k @@ -3488,11 +3542,10 @@ fn render_footer( HintItem::new(key!(Enter), "open"), HintItem::new(key!(Tab), "input"), ]; - if stoppable { - // Pinned so the stop chip always survives compact - // truncation while an agent row is selected. + if show_ctrl_x { hints.push(HintItem::new(stop, stop_label).pinned()); } + ShortcutsBar::new(&hints) .compact(4, Some(HintItem::new(help, "shortcuts"))) .render(inner, buf); @@ -3590,7 +3643,7 @@ fn render_footer( tab_hint, esc_hint, ]; - if stoppable { + if show_ctrl_x { h.push(HintItem::new(stop, stop_label).pinned()); } h @@ -3611,7 +3664,7 @@ fn render_footer( if !reply_empty { h.insert(1, HintItem::new(send_open, "send+open")); } - if stoppable { + if show_ctrl_x { h.push(HintItem::new(stop, stop_label).pinned()); } h @@ -3623,7 +3676,7 @@ fn render_footer( tab_hint, esc_hint, ]; - if stoppable { + if show_ctrl_x { h.push(HintItem::new(stop, stop_label).pinned()); } h @@ -3639,7 +3692,7 @@ fn render_footer( // bare Enter still attaches. let open_key = if peek_focused { send_key } else { enter }; let mut h = vec![HintItem::new(open_key, "open"), tab_hint, esc_hint]; - if stoppable { + if show_ctrl_x { h.push(HintItem::new(stop, stop_label).pinned()); } h @@ -3711,22 +3764,13 @@ fn render_footer( h.push(HintItem::new(send_key, "send")); h.push(HintItem::new(send_open, "send+open")); } - if stoppable { - // Pinned so Ctrl+x always shows while an agent row is - // selected, even if earlier chips would otherwise fill the - // compact bar. + if show_ctrl_x { h.push(HintItem::new(stop, stop_label).pinned()); } h } else { // Defensive — neither the button nor a row is focused. - // Should never happen given the invariant on - // `DashboardState`, but a fall-through keeps the bar - // populated rather than silently empty. - vec![ - HintItem::new(send_key, "create"), - HintItem::new(stop, stop_label), - ] + vec![HintItem::new(send_key, "create")] }; ShortcutsBar::new(&hints) @@ -4437,6 +4481,57 @@ mod tests { ); } + /// Hover `[✗]` paints only on settled rows, never on a busy one. + #[test] + fn render_dashboard_hover_shows_delete_x_only_for_settled_rows() { + use crate::app::roster::{RosterActivity, RosterEntry, RosterOrigin}; + + let ballot = crate::glyphs::ballot_x_button(); + let render_with = |activity: RosterActivity| -> String { + let area = Rect::new(0, 0, 100, 24); + let mut buf = Buffer::empty(area); + let mut agents: IndexMap = IndexMap::new(); + let mut state = DashboardState::new(); + let registry = crate::actions::ActionRegistry::defaults(); + let roster = [RosterEntry { + session_id: "sess-hover".into(), + title: Some("Hover me".into()), + cwd: "/repo/work".into(), + is_worktree: false, + model_id: None, + yolo: false, + activity, + resident: true, + last_change_unix_ms: 1_725_000_000_000, + origin: RosterOrigin::default(), + }]; + state.hovered_row = Some(DashboardRowId::Roster { + session_id: "sess-hover".into(), + }); + let _ = render_dashboard( + &mut buf, + area, + &mut state, + &mut agents, + ®istry, + None, + &roster, + false, + None, + ); + buf_to_text(&buf) + }; + + assert!( + render_with(RosterActivity::Completed).contains(ballot), + "hovering a settled (completed) row must show the [✗] delete affordance", + ); + assert!( + !render_with(RosterActivity::Working).contains(ballot), + "hovering a busy (working) row must NOT show the [✗] delete affordance", + ); + } + /// While the local session roster is still loading the empty body /// shows a loading hint instead of the "no agents" copy. #[test] @@ -5257,12 +5352,12 @@ mod tests { #[test] fn render_row_centers_title_only_content() { let theme = Theme::current(); - let state = DashboardState::new(); + let mut state = DashboardState::new(); // Title-only → centered on the middle line. let row = header_test_row(1, RowState::Idle, "solo"); let mut buf = Buffer::empty(Rect::new(0, 0, 40, 3)); - render_row(&mut buf, Rect::new(0, 0, 40, 3), &theme, &row, &state); + render_row(&mut buf, Rect::new(0, 0, 40, 3), &theme, &row, &mut state); assert_eq!(buf[(4, 1)].symbol(), "s", "title must sit on line 1"); assert_eq!(buf[(4, 0)].symbol(), " ", "line 0 must be padding"); assert_eq!(buf[(4, 2)].symbol(), " ", "line 2 must be padding"); @@ -5271,7 +5366,7 @@ mod tests { let mut row = header_test_row(2, RowState::Working, "pair"); row.secondary_line = Some("Responding".to_string()); let mut buf = Buffer::empty(Rect::new(0, 0, 40, 3)); - render_row(&mut buf, Rect::new(0, 0, 40, 3), &theme, &row, &state); + render_row(&mut buf, Rect::new(0, 0, 40, 3), &theme, &row, &mut state); assert_eq!(buf[(4, 0)].symbol(), "p", "title must sit on line 0"); assert_eq!(buf[(4, 1)].symbol(), "R", "secondary must sit on line 1"); assert_eq!(buf[(4, 2)].symbol(), " ", "line 2 must be padding"); @@ -6787,7 +6882,7 @@ mod tests { is_more_placeholder: false, more_count: 0, }; - render_row(&mut buf, Rect::new(0, 0, 100, 2), &theme, &row, &state); + render_row(&mut buf, Rect::new(0, 0, 100, 2), &theme, &row, &mut state); // Title row. assert_eq!( @@ -6877,13 +6972,13 @@ mod tests { // Unselected → dim secondary. let mut buf = Buffer::empty(Rect::new(0, 0, 100, 2)); - let state_unselected = DashboardState::new(); + let mut state_unselected = DashboardState::new(); render_row( &mut buf, Rect::new(0, 0, 100, 2), &theme, &row, - &state_unselected, + &mut state_unselected, ); assert_eq!( buf[(4, 1)].fg, @@ -6900,7 +6995,7 @@ mod tests { Rect::new(0, 0, 100, 2), &theme, &row, - &state_selected, + &mut state_selected, ); assert_eq!( buf[(4, 1)].fg, @@ -6946,7 +7041,7 @@ mod tests { Rect::new(0, 0, 100, 2), &theme, &make_row(), - &state, + &mut state, ); buf }; @@ -7007,7 +7102,7 @@ mod tests { use std::time::SystemTime; let mut buf = Buffer::empty(Rect::new(0, 0, 100, 2)); let theme = Theme::current(); - let state = DashboardState::new(); + let mut state = DashboardState::new(); let row = DashboardRow { id: DashboardRowId::TopLevel(crate::app::agent::AgentId(1)), label: "New session #abc12345".to_string(), @@ -7027,7 +7122,7 @@ mod tests { is_more_placeholder: false, more_count: 0, }; - render_row(&mut buf, Rect::new(0, 0, 100, 2), &theme, &row, &state); + render_row(&mut buf, Rect::new(0, 0, 100, 2), &theme, &row, &mut state); // Title starts at col 4: "New session" (11 chars, cols 4..15) then // " #abc12345" (suffix from col 15). @@ -8030,10 +8125,6 @@ mod tests { ); } - /// When peek is active, the footer flips to peek- - /// mode hints (`enter:open · esc:New Agent · ctrl+x:close`). The - /// nav chip is dropped (saving space) and the Ctrl+x chip stays - /// visible while the agent is selected. #[test] fn render_footer_peek_mode_shows_peek_hints() { let mut buf = Buffer::empty(Rect::new(0, 0, 200, 1)); @@ -8046,8 +8137,8 @@ mod tests { &theme, &state, ®istry, - None, - true, // peek_active + Some(RowState::Idle), + true, None, ); let content = buf_to_text(&buf); @@ -8055,11 +8146,9 @@ mod tests { content.contains(":open") && content.contains(":New Agent"), "peek-mode footer must include open + New Agent (unselect) hints, got: {content:?}", ); - // The stop chip stays visible while an agent is selected. With - // no row state passed (None) the label is the idle-style `close`. assert!( - content.contains(":close"), - "peek-mode footer must keep the Ctrl+x stop chip, got: {content:?}", + content.contains(":delete"), + "peek-mode footer must keep the Ctrl+x delete chip, got: {content:?}", ); // The nav chip is dropped to save bottom-bar space. assert!( @@ -8373,10 +8462,8 @@ mod tests { ); } - /// An Inactive (roster-only) selection has nothing running to stop — - /// the stop chip is suppressed in both focus modes. #[test] - fn render_footer_inactive_row_hides_stop() { + fn render_footer_inactive_row_shows_delete() { let theme = Theme::current(); let registry = crate::actions::ActionRegistry::defaults(); @@ -8396,15 +8483,14 @@ mod tests { ); let content = buf_to_text(&buf); assert!( - !content.contains(":stop") && !content.contains(":close"), - "inactive row footer must NOT show the stop chip, got: {content:?}", + content.contains(":delete"), + "list-focused idle-row footer must show the delete chip, got: {content:?}", ); assert!( content.contains(":open"), - "inactive row footer keeps the open chip, got: {content:?}", + "list-focused idle-row footer must show the open chip, got: {content:?}", ); - // List focused (Tab) — same suppression. state.list_focused = true; let mut buf2 = Buffer::empty(Rect::new(0, 0, 200, 1)); render_footer( @@ -8418,13 +8504,8 @@ mod tests { None, ); let content2 = buf_to_text(&buf2); - assert!( - !content2.contains(":stop") && !content2.contains(":close"), - "list-focused inactive footer must NOT show the stop chip, got: {content2:?}", - ); + assert!(content2.contains(":delete"), "{content2:?}"); - // Control: an Idle selection keeps the stop chip in both modes — - // labelled `close` (the session is idle, so Ctrl+x closes it). state.list_focused = false; let mut buf3 = Buffer::empty(Rect::new(0, 0, 200, 1)); render_footer( @@ -8438,16 +8519,9 @@ mod tests { None, ); let content3 = buf_to_text(&buf3); - assert!( - content3.contains(":close"), - "idle row footer must keep the stop chip labelled `close`, got: {content3:?}", - ); + assert!(content3.contains(":delete"), "{content3:?}"); } - /// The Ctrl+x chip label follows the selected agent's state: a - /// Working or NeedsInput agent shows `stop` (cancel the turn — a - /// NeedsInput row keeps a paused-but-running turn), while an idle / - /// quiet one shows `close` (close the session). #[test] fn render_footer_stop_label_follows_state() { let theme = Theme::current(); @@ -8492,7 +8566,7 @@ mod tests { "NeedsInput agent footer must label Ctrl+x as `stop`, got: {needs_input:?}", ); - // Idle → `close`. + // Idle → `delete`. let mut buf2 = Buffer::empty(Rect::new(0, 0, 200, 1)); render_footer( &mut buf2, @@ -8506,8 +8580,8 @@ mod tests { ); let idle = buf_to_text(&buf2); assert!( - idle.contains(":close") && !idle.contains(":stop"), - "Idle agent footer must label Ctrl+x as `close`, got: {idle:?}", + idle.contains(":delete") && !idle.contains(":stop"), + "Idle agent footer must label Ctrl+x as `delete`, got: {idle:?}", ); } @@ -8832,17 +8906,14 @@ mod tests { ); } - /// Stop-confirm armed routes through `ShortcutsBar::with_pending`. + /// Delete-confirm armed while the input is focused routes through + /// `ShortcutsBar::with_pending` ("press Ctrl+x again to delete"). #[test] - fn render_footer_stop_confirm_uses_pending_hint() { - use std::time::Instant; + fn render_footer_delete_confirm_uses_pending_hint() { let mut buf = Buffer::empty(Rect::new(0, 0, 200, 1)); let theme = Theme::current(); let mut state = DashboardState::new(); - state.stop_confirm = Some(( - DashboardRowId::TopLevel(crate::app::agent::AgentId(1)), - Instant::now(), - )); + state.arm_delete(DashboardRowId::TopLevel(crate::app::agent::AgentId(1))); let registry = crate::actions::ActionRegistry::defaults(); render_footer( &mut buf, @@ -8860,26 +8931,26 @@ mod tests { "stop-confirm footer must say `press again`, got: {content:?}", ); assert!( - content.to_lowercase().contains("close this session"), - "stop-confirm footer must mention closing the session, got: {content:?}", + content.to_lowercase().contains("delete this session"), + "delete-confirm footer must name the action, got: {content:?}", ); } - /// An EXPIRED stop-confirm (older than `STOP_CONFIRM_WINDOW`) must + /// An EXPIRED delete-confirm (older than `CONFIRM_WINDOW`) must /// not claim the footer — the dispatcher would re-arm rather than - /// close on the next press, so "press again" would lie. Regular + /// delete on the next press, so "press again" would lie. Regular /// hints render instead (e.g. after a mouse click moved the /// selection without a keypress to disarm the confirm). #[test] - fn render_footer_expired_stop_confirm_shows_regular_hints() { + fn render_footer_expired_delete_confirm_shows_regular_hints() { use std::time::{Duration, Instant}; let mut buf = Buffer::empty(Rect::new(0, 0, 200, 1)); let theme = Theme::current(); let mut state = DashboardState::new(); state.focus_row(DashboardRowId::TopLevel(crate::app::agent::AgentId(1))); - state.stop_confirm = Some(( + state.delete_confirm = Some(( DashboardRowId::TopLevel(crate::app::agent::AgentId(1)), - Instant::now() - (super::super::state::STOP_CONFIRM_WINDOW + Duration::from_secs(1)), + Instant::now() - (super::super::state::CONFIRM_WINDOW + Duration::from_secs(1)), )); let registry = crate::actions::ActionRegistry::defaults(); render_footer( diff --git a/crates/codegen/xai-grok-pager/src/views/dashboard/row.rs b/crates/codegen/xai-grok-pager/src/views/dashboard/row.rs index f39a353..7b30131 100644 --- a/crates/codegen/xai-grok-pager/src/views/dashboard/row.rs +++ b/crates/codegen/xai-grok-pager/src/views/dashboard/row.rs @@ -253,7 +253,9 @@ fn build_local_rows( rows } /// Map a leader [`RosterActivity`] to the dashboard's coarse [`RowState`]. -fn roster_activity_to_state(activity: RosterActivity) -> RowState { +/// Public so the dispatcher can gate roster-row deletion through the very +/// same `RowState::allows_delete` predicate the renderer paints `[✗]` with. +pub fn roster_activity_to_state(activity: RosterActivity) -> RowState { match activity { RosterActivity::Working => RowState::Working, RosterActivity::NeedsInput => RowState::NeedsInput, diff --git a/crates/codegen/xai-grok-pager/src/views/dashboard/state.rs b/crates/codegen/xai-grok-pager/src/views/dashboard/state.rs index 8d15c70..7c4c8d5 100644 --- a/crates/codegen/xai-grok-pager/src/views/dashboard/state.rs +++ b/crates/codegen/xai-grok-pager/src/views/dashboard/state.rs @@ -181,11 +181,10 @@ impl PersistedRowId { } } -/// Window within which a second `Ctrl+X` press confirms closing the -/// selected agent. Shared by the dispatcher (which gates the actual -/// close) and the footer (which only paints the "press again" hint -/// while the window is live). -pub const STOP_CONFIRM_WINDOW: std::time::Duration = std::time::Duration::from_secs(2); +/// Window within which a second confirming gesture (`Ctrl+X`, a `[✗]` +/// click, or `y`) deletes the armed row. Also reused by the +/// dashboard-overlay stop for its double-press close confirm. +pub const CONFIRM_WINDOW: std::time::Duration = std::time::Duration::from_secs(2); /// Coarse state used for the dashboard grouping. /// @@ -215,6 +214,17 @@ pub enum RowState { } impl RowState { + /// The one predicate for "may be deleted", shared by the renderer's + /// `[✗]` and the dispatcher: only settled rows qualify. `Working` / + /// `NeedsInput` are excluded so an in-flight turn is never wiped — + /// `Ctrl+X` cancels those instead. + pub fn allows_delete(self) -> bool { + matches!( + self, + Self::Idle | Self::Inactive | Self::Completed | Self::Failed + ) + } + /// Sort priority used inside a state group: higher = floats up. /// Pinned rows always float to the absolute top regardless of state. pub fn group_priority(self) -> u8 { @@ -493,11 +503,10 @@ pub struct DashboardState { /// exists"). Rendered verbatim by `paint_dispatch_feedback_badge`; /// error messages are built via [`Self::set_error_toast`]. pub error_toast: Option, - /// Pending stop confirmation. `Some((row, set_at))` after the first - /// `Ctrl+X` press on a top-level row. The second press within - /// [`STOP_CONFIRM_WINDOW`] closes the agent. Mirrors the session-close - /// close-confirm pattern. - pub stop_confirm: Option<(DashboardRowId, Instant)>, + /// Row armed for delete, and when. A second gesture on the same row + /// within [`CONFIRM_WINDOW`] deletes it (see [`Self::armed_delete_row`]); + /// otherwise it lapses. Cleared on any focus change. + pub delete_confirm: Option<(DashboardRowId, Instant)>, /// Tick counter for spinner animation. The /// counter is bumped by [`crate::app::app_view::AppView::tick`] /// (NOT the renderer, which is read-only). @@ -508,6 +517,15 @@ pub struct DashboardState { /// mouse handling to map (col, row) → row id without scanning the /// row list a second time. pub row_rects: Vec<(DashboardRowId, Rect)>, + /// Per-row `[✗]` hit areas, rebuilt each render; maps a click onto the + /// delete gesture instead of a row select. + pub row_delete_rects: Vec<(DashboardRowId, Rect)>, + /// Row whose `[✗]` the mouse is over, so the renderer can tint it. + pub hovered_delete: Option, + /// Roster session ids whose origin is a chat `conversation` — those + /// can't be deleted from the dashboard yet, so they get no `[✗]` and + /// don't arm. Rebuilt each render from the roster. + pub conversation_row_ids: std::collections::HashSet, /// Last frame's section-header hit areas keyed by [`SectionKey`]. /// Used by mouse handling to map (col, row) → section for /// click-to-toggle and hover. Rebuilt every render. @@ -1351,9 +1369,12 @@ impl DashboardState { peek_reply_target_cwd: None, rename: None, error_toast: None, - stop_confirm: None, + delete_confirm: None, spinner_tick: 0, row_rects: Vec::new(), + row_delete_rects: Vec::new(), + hovered_delete: None, + conversation_row_ids: std::collections::HashSet::new(), section_rects: Vec::new(), idle_overflow_rect: None, last_area: Rect::default(), @@ -1474,6 +1495,7 @@ impl DashboardState { self.selected = None; self.selected_section = None; self.selected_idle_overflow = false; + self.delete_confirm = None; } /// Focus the row identified by `id`. Clears the @@ -1483,6 +1505,13 @@ impl DashboardState { /// risk — the invariant only holds when both fields are /// written through here. pub fn focus_row(&mut self, id: DashboardRowId) { + if self + .delete_confirm + .as_ref() + .is_some_and(|(armed, _)| armed != &id) + { + self.delete_confirm = None; + } self.selected = Some(id); self.new_agent_button_focused = false; self.selected_section = None; @@ -1497,6 +1526,7 @@ impl DashboardState { self.selected = None; self.new_agent_button_focused = false; self.selected_idle_overflow = false; + self.delete_confirm = None; } /// Focus the Idle group's "N more" overflow toggle — @@ -1507,6 +1537,62 @@ impl DashboardState { self.selected = None; self.selected_section = None; self.new_agent_button_focused = false; + self.delete_confirm = None; + } + + fn set_list_focused(&mut self, focused: bool) { + self.list_focused = focused; + if !focused { + self.delete_confirm = None; + } + } + + /// The armed row while its [`CONFIRM_WINDOW`] is still live, clearing + /// an expired arm as a side effect. The accessor the dispatcher and + /// mouse handler share so "armed on screen" and "armed for delete" + /// never diverge. + pub fn armed_delete_row(&mut self) -> Option { + match &self.delete_confirm { + Some((id, at)) if at.elapsed() < CONFIRM_WINDOW => Some(id.clone()), + Some(_) => { + self.delete_confirm = None; + None + } + None => None, + } + } + + /// Read-only counterpart of [`Self::armed_delete_row`] for the + /// renderer (does not clear an expired arm). + pub fn armed_delete_row_ref(&self) -> Option<&DashboardRowId> { + self.delete_confirm + .as_ref() + .filter(|(_, at)| at.elapsed() < CONFIRM_WINDOW) + .map(|(id, _)| id) + } + + pub fn arm_delete(&mut self, id: DashboardRowId) { + self.delete_confirm = Some((id, Instant::now())); + } + + /// Whether `id` is a chat-conversation roster row, which the dashboard + /// can't delete yet (see [`Self::conversation_row_ids`]). + pub fn row_is_conversation(&self, id: &DashboardRowId) -> bool { + matches!(id, DashboardRowId::Roster { session_id } + if self.conversation_row_ids.contains(session_id)) + } + + /// Enforce the invariant that a delete arm belongs to the selected + /// row. Selection changes routed through the focus helpers already + /// disarm, but `reanchor_selection` / `gc_stale_refs` can drop or move + /// `selected` directly — without this a stale arm would let a later + /// `y` delete a row that is no longer selected. + fn sync_delete_confirm_to_selection(&mut self) { + if let Some((armed, _)) = self.delete_confirm.as_ref() + && self.selected.as_ref() != Some(armed) + { + self.delete_confirm = None; + } } /// Toggle whether the Idle group shows every agent (`true`) or caps @@ -1683,6 +1769,7 @@ impl DashboardState { // holds at every close site, not just here. self.close_popup(); } + self.sync_delete_confirm_to_selection(); } /// Switch grouping (`Ctrl+G`). @@ -3001,8 +3088,37 @@ impl DashboardState { InputOutcome::Action(Action::DashboardDispatch { text, attach }) } + /// List-focused `y`/`n` confirm for an already-armed delete (arming is + /// via `Ctrl+X` / `[✗]`, not `d`). When the list isn't focused, + /// disarming is left to the caller so a second `Ctrl+X` reaches the + /// dispatcher. + fn handle_delete_confirm_key(&mut self, key: &KeyEvent) -> Option { + if key.kind == KeyEventKind::Release { + return None; + } + if !self.list_focused { + return None; + } + self.armed_delete_row()?; + if !key.modifiers.is_empty() { + self.delete_confirm = None; + return None; + } + match key.code { + KeyCode::Char('y') => Some(InputOutcome::Action(Action::DashboardDelete)), + KeyCode::Char('n') => { + self.delete_confirm = None; + Some(InputOutcome::Changed) + } + _ => { + self.delete_confirm = None; + None + } + } + } + fn handle_key(&mut self, key: &KeyEvent, registry: &ActionRegistry) -> InputOutcome { - // Resolve the registry binding up-front — the toast / stop-confirm + // Resolve the registry binding up-front — the toast / delete-confirm // clear below needs to know whether this key IS the stop key, and // it must run before the peek intercept (the lookup itself is a // pure read; the action is honoured further down). @@ -3017,37 +3133,26 @@ impl DashboardState { let from_registry = registry.lookup_with_mode(key, crate::actions::When::DashboardFocused, vim_mode); - // Clear `error_toast` at the TOP of the - // handler so any subsequent keypress dismisses the toast, - // regardless of which branch handles the key (including keys - // the peek panel consumes — peek is open by default for a - // selected row, so nav keys route through it). - // - // When the toast is cleared, the linked - // `stop_confirm` armed state is also cleared. The two state - // bits are semantically linked: the user saw "Press Ctrl+X - // again", that hint is now gone, so re-arm rather than let a - // stale confirm window silently close the wrong session. - // - // The clear is SKIPPED when the resolved - // action is `DashboardStop`. Without this skip, the second - // Ctrl+X press would wipe the just-armed `stop_confirm` - // before `dispatch_dashboard_stop` could observe it, and the - // session would never close (the dispatcher kept re-arming a - // fresh confirm on every press). The Ctrl+X path owns - // `stop_confirm` and `error_toast` end-to-end: the first - // press arms both, the second press observes them and closes. - let preserve_stop_state = - matches!(from_registry, Some(crate::actions::ActionId::DashboardStop)); - if !preserve_stop_state { + // Clear `error_toast` on any keypress so it never lingers; kept for + // `Ctrl+X` so the arm path's own messaging survives its first press. + let is_stop_key = matches!(from_registry, Some(crate::actions::ActionId::DashboardStop)); + if !is_stop_key { self.error_toast = None; - // The disarm is NOT gated on `error_toast` being set (the - // Ctrl+X arm path deliberately plants no toast): a pending - // stop confirmation is bound to the row that was selected - // when Ctrl+X was pressed, so any other key — nav included — - // must disarm it. Otherwise the footer's "press again to - // close" hint lingers while the cursor moves to other agents. - self.stop_confirm = None; + } + + // Disarm delete-confirm on any non-confirming key. Two gestures are + // preserved: `Ctrl+X` (its second press is the confirm, read by the + // dispatcher) and a list-focused bare `y`/`n` (handled just below). + let confirm_via_yn = self.list_focused + && self.armed_delete_row().is_some() + && key.modifiers.is_empty() + && matches!(key.code, KeyCode::Char('y') | KeyCode::Char('n')); + if !is_stop_key && !confirm_via_yn { + self.delete_confirm = None; + } + + if !is_stop_key && let Some(outcome) = self.handle_delete_confirm_key(key) { + return outcome; } // Free-tier override: Ctrl+O opens the pinned upgrade CTA (when one is @@ -3405,6 +3510,13 @@ impl DashboardState { } _ => true, }; + // Never let an auto-repeat (held key) drive the destructive + // Ctrl+X arm→confirm — holding the key would arm and immediately + // confirm a delete. Require discrete presses, like the picker's + // `y` confirm. Non-destructive actions may still repeat. + if id == crate::actions::ActionId::DashboardStop && key.kind == KeyEventKind::Repeat { + return InputOutcome::Unchanged; + } if honor && let Some(outcome) = dashboard_action_for_id(id, &mut self.error_toast) { return outcome; } @@ -3483,7 +3595,7 @@ impl DashboardState { // slash / `@` dropdowns are open the intercepts above already // consumed Tab (accept completion), so this only fires otherwise. if matches!(key.code, KeyCode::Tab) && key.modifiers.is_empty() { - self.list_focused = !self.list_focused; + self.set_list_focused(!self.list_focused); // Re-engage selection-follow so the viewport tracks the // cursor once the list takes focus. self.clear_manual_scroll(); @@ -3502,12 +3614,12 @@ impl DashboardState { { if vim_mode { if key.code == KeyCode::Char('i') && key.modifiers.is_empty() { - self.list_focused = false; + self.set_list_focused(false); return InputOutcome::Changed; } return InputOutcome::Unchanged; } - self.list_focused = false; + self.set_list_focused(false); // fall through to the widget so the char is typed. } else { // Non-printable (Backspace/Home/…) while the overview is @@ -3636,6 +3748,20 @@ impl DashboardState { self.hovered_row = new_hover; changed = true; } + let new_hover_delete = self + .row_delete_rects + .iter() + .find(|(_, r)| { + mouse.column >= r.x + && mouse.column < r.x + r.width + && mouse.row >= r.y + && mouse.row < r.y + r.height + }) + .map(|(id, _)| id.clone()); + if new_hover_delete != self.hovered_delete { + self.hovered_delete = new_hover_delete; + changed = true; + } // Section-header hover → the renderer brightens its text. let new_hover_section = self .section_rects @@ -3773,7 +3899,7 @@ impl DashboardState { self.dispatch.accept_slash_completion(&self.models); } } - self.list_focused = false; + self.set_list_focused(false); return InputOutcome::Changed; } @@ -3821,7 +3947,29 @@ impl DashboardState { } } } - self.list_focused = false; + self.set_list_focused(false); + return InputOutcome::Changed; + } + + if let Some(id) = self + .row_delete_rects + .iter() + .find(|(_, r)| { + mouse.column >= r.x + && mouse.column < r.x + r.width + && mouse.row >= r.y + && mouse.row < r.y + r.height + }) + .map(|(id, _)| id.clone()) + { + self.manual_scroll_active = false; + // Second `[✗]` click within the window confirms; else re-arm. + if self.armed_delete_row().as_ref() == Some(&id) { + return InputOutcome::Action(Action::DashboardDelete); + } + self.focus_row(id.clone()); + self.set_list_focused(true); + self.arm_delete(id); return InputOutcome::Changed; } @@ -3934,7 +4082,7 @@ impl DashboardState { && mouse.row >= rect.y && mouse.row < rect.y + rect.height { - self.list_focused = false; + self.set_list_focused(false); // Forward the click so the caret lands where the user // clicked. Skipped in search mode, where the prompt // renders its own single-line cursor with a `Search:` @@ -4334,6 +4482,7 @@ impl DashboardState { rows.iter().filter(|r| !r.is_more_placeholder).collect(); if selectable.is_empty() { self.selected = None; + self.delete_confirm = None; return; } if let Some(sel) = self.selected.as_ref() @@ -4344,6 +4493,7 @@ impl DashboardState { // the user's job. self.selected = None; } + self.sync_delete_confirm_to_selection(); } } @@ -9586,33 +9736,33 @@ mod tests { ); } - /// An armed stop confirmation is bound to the row that was selected + /// An armed delete confirmation is bound to the row that was selected /// when `Ctrl+X` was pressed — any other key (nav included) must - /// disarm it, otherwise the footer's "press again to close" hint + /// disarm it, otherwise the footer's "press again to delete" hint /// lingers while the cursor moves to other agents. The disarm must /// NOT depend on `error_toast` (the Ctrl+X arm path plants none). #[test] - fn nav_key_disarms_pending_stop_confirm() { + fn nav_key_disarms_pending_delete_confirm() { let mut state = DashboardState::new(); let reg = crate::actions::ActionRegistry::defaults(); state.focus_row(DashboardRowId::TopLevel(AgentId(0))); - state.stop_confirm = Some((DashboardRowId::TopLevel(AgentId(0)), Instant::now())); + state.arm_delete(DashboardRowId::TopLevel(AgentId(0))); assert!(state.error_toast.is_none(), "arm path plants no toast"); let _ = state.handle_key(&KeyEvent::new(KeyCode::Down, KeyModifiers::NONE), ®); assert!( - state.stop_confirm.is_none(), - "a nav keypress must disarm the pending stop confirm", + state.delete_confirm.is_none(), + "a nav keypress must disarm the pending delete confirm", ); // Control — Ctrl+X itself preserves the armed confirm so the - // dispatcher can observe it and close. - state.stop_confirm = Some((DashboardRowId::TopLevel(AgentId(0)), Instant::now())); + // dispatcher can observe it and delete. + state.arm_delete(DashboardRowId::TopLevel(AgentId(0))); let _ = state.handle_key( &KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL), ®, ); assert!( - state.stop_confirm.is_some(), + state.delete_confirm.is_some(), "Ctrl+X must preserve the armed confirm for the dispatcher", ); @@ -9620,18 +9770,117 @@ mod tests { // row, and `handle_peek_key` CONSUMES Up/Down (agent switch) — // the disarm must sit above that intercept or nav keys never // reach it and the footer hint lingers. - state.stop_confirm = Some((DashboardRowId::TopLevel(AgentId(0)), Instant::now())); + state.arm_delete(DashboardRowId::TopLevel(AgentId(0))); state.peek = Some(super::super::peek::PeekPanelState::new( DashboardRowId::TopLevel(AgentId(0)), peek_fields_for_test("Idle"), )); let _ = state.handle_key(&KeyEvent::new(KeyCode::Down, KeyModifiers::NONE), ®); assert!( - state.stop_confirm.is_none(), + state.delete_confirm.is_none(), "a nav keypress consumed by the peek panel must still disarm the confirm", ); } + #[test] + fn click_delete_control_arms_then_confirms() { + use crossterm::event::{MouseButton, MouseEvent, MouseEventKind}; + let mut state = DashboardState::new(); + let id = DashboardRowId::TopLevel(AgentId(0)); + state + .row_delete_rects + .push((id.clone(), Rect::new(10, 2, 3, 1))); + let click = |col, row| MouseEvent { + kind: MouseEventKind::Down(MouseButton::Left), + column: col, + row, + modifiers: KeyModifiers::NONE, + }; + // First `[✗]` click only arms — it must not open/attach the session. + let first = state.handle_mouse(&click(11, 2)); + assert!(matches!(first, InputOutcome::Changed), "got {first:?}"); + assert!(!matches!( + first, + InputOutcome::Action(Action::DashboardAttach(_)) + )); + assert_eq!(state.armed_delete_row_ref(), Some(&id)); + // Second click confirms. + assert!(matches!( + state.handle_mouse(&click(11, 2)), + InputOutcome::Action(Action::DashboardDelete) + )); + } + + #[test] + fn focus_change_disarms_delete_confirm() { + let mut state = DashboardState::new(); + let a = DashboardRowId::TopLevel(AgentId(0)); + let b = DashboardRowId::TopLevel(AgentId(1)); + state.focus_row(a.clone()); + state.arm_delete(a.clone()); + state.focus_row(a.clone()); + assert_eq!(state.armed_delete_row_ref(), Some(&a)); + state.focus_row(b); + assert!(state.delete_confirm.is_none()); + state.arm_delete(DashboardRowId::TopLevel(AgentId(0))); + state.focus_new_agent_button(); + assert!(state.delete_confirm.is_none()); + + state.focus_row(a.clone()); + state.list_focused = true; + state.arm_delete(a); + state.dispatch_rect = Some(Rect::new(0, 10, 40, 1)); + let _ = state.handle_mouse(&crossterm::event::MouseEvent { + kind: crossterm::event::MouseEventKind::Down(crossterm::event::MouseButton::Left), + column: 2, + row: 10, + modifiers: KeyModifiers::NONE, + }); + assert!(state.delete_confirm.is_none()); + assert!(!state.list_focused); + } + + /// An auto-repeat (held) Ctrl+X must not drive the destructive + /// arm→confirm: only discrete presses count, so holding the key can't + /// arm and immediately confirm a delete. + #[test] + fn ctrl_x_key_repeat_is_ignored() { + let mut state = DashboardState::new(); + let reg = crate::actions::ActionRegistry::defaults(); + state.focus_row(DashboardRowId::TopLevel(AgentId(0))); + let repeat = Event::Key(crossterm::event::KeyEvent { + code: KeyCode::Char('x'), + modifiers: KeyModifiers::CONTROL, + kind: crossterm::event::KeyEventKind::Repeat, + state: crossterm::event::KeyEventState::NONE, + }); + assert!(matches!( + state.handle_input(&repeat, ®), + InputOutcome::Unchanged + )); + // A real press still resolves to the stop action. + let press = Event::Key(KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL)); + assert!(matches!( + state.handle_input(&press, ®), + InputOutcome::Action(Action::DashboardStop) + )); + } + + /// `gc_stale_refs` dropping the selected row (session left the list) + /// must also disarm delete, so a later `y` can't delete a phantom row. + #[test] + fn gc_stale_refs_disarms_delete_when_selection_dropped() { + let mut state = DashboardState::new(); + let a = DashboardRowId::TopLevel(AgentId(0)); + state.focus_row(a.clone()); + state.arm_delete(a.clone()); + assert!(state.armed_delete_row_ref().is_some()); + // The armed row is no longer alive → gc drops selection AND disarms. + state.gc_stale_refs(&|_| false); + assert!(state.selected.is_none()); + assert!(state.delete_confirm.is_none(), "stale arm must be cleared"); + } + /// Section header selected while the LIST is focused — the input is /// inactive, so Enter / Left / Right operate on the section even /// when a draft is sitting in the (unfocused) dispatch input. diff --git a/crates/codegen/xai-grok-pager/src/views/modal.rs b/crates/codegen/xai-grok-pager/src/views/modal.rs index a9312f4..7a177c7 100644 --- a/crates/codegen/xai-grok-pager/src/views/modal.rs +++ b/crates/codegen/xai-grok-pager/src/views/modal.rs @@ -229,11 +229,9 @@ pub enum ActiveModal { entries_query: Option, /// Source filter for the modal session picker. source_filter: crate::views::session_picker::SourceFilter, - /// Session armed for delete, captured as `(source, session_id, cwd)` when - /// `d` is pressed so the `y` confirm always has a valid cwd even if - /// the picker lists change underneath it. `Some` only while the - /// focused row is armed; cleared on cancel / completion. - pending_delete: Option<(String, String, String)>, + /// Session armed for delete via `d` (see + /// [`crate::views::session_picker::PendingDelete`]). + pending_delete: Option, }, /// How-to documentation list modal (wider picker style). DocPicker { diff --git a/crates/codegen/xai-grok-pager/src/views/question_view.rs b/crates/codegen/xai-grok-pager/src/views/question_view.rs index 7e4f3cc..f48318a 100644 --- a/crates/codegen/xai-grok-pager/src/views/question_view.rs +++ b/crates/codegen/xai-grok-pager/src/views/question_view.rs @@ -201,6 +201,11 @@ pub struct QuestionViewState { /// while the user is answering questions — the time spent in the /// question view is subtracted from the turn elapsed display. pub opened_at: Instant, + /// Wall-clock twin of `opened_at` (UTC ms). `Instant` is suspend-blind, + /// so a pause netted against the wall-anchored turn span must itself be + /// measured on the wall clock, or a suspend during an open question + /// would read as worked time. + pub opened_at_wall_ms: i64, /// When `true`, the freeform "Other" input row is hidden. Used by /// locally-driven questions (e.g. credit-limit upsell) that only /// offer fixed options with no free-text fallback. @@ -272,6 +277,7 @@ impl QuestionViewState { bottom_panel_index: None, local_kind: None, opened_at: Instant::now(), + opened_at_wall_ms: chrono::Utc::now().timestamp_millis(), no_freeform: false, } } diff --git a/crates/codegen/xai-grok-pager/src/views/session_picker.rs b/crates/codegen/xai-grok-pager/src/views/session_picker.rs index 3279ef4..b101007 100644 --- a/crates/codegen/xai-grok-pager/src/views/session_picker.rs +++ b/crates/codegen/xai-grok-pager/src/views/session_picker.rs @@ -78,6 +78,89 @@ pub enum PickerItem { Content { hit_index: usize }, } +/// A session armed for deletion, captured on `d` so the `y` confirm keeps +/// a valid `(source, session_id, cwd)` even if the lists shift. Shared by +/// the welcome and modal `/resume` pickers so they can't drift apart. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PendingDelete { + pub source: String, + pub session_id: String, + pub cwd: String, +} + +/// Outcome of routing a key through an armed [`PendingDelete`] confirm. +pub(crate) enum PendingDeleteKey { + /// `y`: caller should delete this session. + Confirm(PendingDelete), + /// `n`: arm cleared; caller should redraw. + Cancel, + /// Other key: arm cleared, but the key should still be processed. + Disarmed, + /// Nothing armed, or not an unmodified key press. + NotArmed, +} + +/// Arm a [`PendingDelete`] from the selected row, or `None` if it can't be +/// deleted (foreign source or non-selectable position). +pub(crate) fn pending_delete_from_selection( + selected: usize, + entry_map: &[Option], + entries: Option<&[SessionPickerEntry]>, + content_results: Option<&[xai_grok_shell::extensions::session_search::SearchSessionHit]>, +) -> Option { + match entry_map.get(selected).and_then(|e| e.as_ref())? { + PickerItem::Fuzzy { original_index } => entries + .and_then(|e| e.get(*original_index)) + .filter(|entry| !crate::app::is_foreign_picker_source(&entry.source)) + .map(|e| PendingDelete { + source: e.source.clone(), + session_id: e.id.clone(), + cwd: e.cwd.clone(), + }), + PickerItem::Content { hit_index } => { + content_results + .and_then(|h| h.get(*hit_index)) + .map(|h| PendingDelete { + source: "local".into(), + session_id: h.session_id.clone(), + cwd: h.cwd.clone(), + }) + } + } +} + +/// Route a key through an armed [`PendingDelete`]: `y` confirms, `n` +/// cancels, any other unmodified key disarms and falls through. +pub(crate) fn handle_pending_delete_key( + pending: &mut Option, + ev: &crossterm::event::Event, +) -> PendingDeleteKey { + use crossterm::event::{Event, KeyCode, KeyEventKind}; + if pending.is_none() { + return PendingDeleteKey::NotArmed; + } + let Event::Key(k) = ev else { + return PendingDeleteKey::NotArmed; + }; + if k.kind != KeyEventKind::Press || !k.modifiers.is_empty() { + return PendingDeleteKey::NotArmed; + } + match k.code { + KeyCode::Char('y') => pending + .take() + .map(PendingDeleteKey::Confirm) + .unwrap_or(PendingDeleteKey::Cancel), + KeyCode::Char('n') => { + *pending = None; + PendingDeleteKey::Cancel + } + _ => { + *pending = None; + PendingDeleteKey::Disarmed + } + } +} + /// Owned data for a single session picker row. Built once per frame and /// then borrowed by `PickerEntry` / `PickerField` slices. Shared between /// the welcome-screen `render_session_picker` and the diff --git a/crates/codegen/xai-grok-pager/src/views/shortcuts_help.rs b/crates/codegen/xai-grok-pager/src/views/shortcuts_help.rs index 5873a8e..4ace5ea 100644 --- a/crates/codegen/xai-grok-pager/src/views/shortcuts_help.rs +++ b/crates/codegen/xai-grok-pager/src/views/shortcuts_help.rs @@ -121,6 +121,23 @@ const REDO_LONG_HELP: &str = "\ Redoes the last undone change in the prompt editor.\n\ Ctrl+Shift+Z is primary; Ctrl+R is an alternate."; +// Prompt history is not an ActionRegistry entry: Up is an inline key handler and +// /history is a slash command. Surface both here for discoverability. +const HISTORY_LONG_HELP: &str = "\ +Recalls previously sent prompts.\n\ +Press Up on an empty prompt to browse earlier prompts, newest first; each move \ +live-populates the composer so you can edit and resend.\n\ +Run /history to open a searchable history panel and filter by text."; + +// Scrollback search has no ActionRegistry entry: it's the vim `/` inline handler, +// or the /find slash command in simple mode. Surface both triggers here. +const SCROLLBACK_SEARCH_LONG_HELP: &str = "\ +Searches the conversation scrollback for text and jumps between matches.\n\ +In the prompt input, run /find to search. In vim mode, you can also press / \ +while the scrollback is focused.\n\ +Type a query, then use n and N (or the arrow keys) to step through matches. \ +Press Enter to jump to a match and Esc to dismiss."; + /// Build the entries vector for the modal, grouped by category. /// /// All registered actions are included, grouped by category. Actions @@ -273,7 +290,25 @@ pub fn build_entries( item, dimmed, action_id: None, - long_help: None, + long_help: Some(SCROLLBACK_SEARCH_LONG_HELP), + }); + } + // Simple mode reaches scrollback search via the `/find` slash command, + // not a keystroke: use a null key + custom display so the raw key list + // stays empty of `/`. + if !vim_mode && cat == Category::ConversationNav { + let mut item = HintItem::new(crate::key!(Null), "search"); + item.custom_display = Some("/find"); + item.description = Some("Search scrollback".into()); + // `/find` is a slash command typed at the prompt (not a scrollback + // keystroke like the vim `/` above), so it is available when the + // prompt is focused — dim on `!PromptFocused`, not scrollback. + let dimmed = !active_contexts.contains(&When::PromptFocused); + entries.push(ShortcutsHelpEntry::Hint { + item, + dimmed, + action_id: None, + long_help: Some(SCROLLBACK_SEARCH_LONG_HELP), }); } // Clipboard + textarea chords not in ActionRegistry. Super/Cmd omitted @@ -308,6 +343,19 @@ pub fn build_entries( redo.description = Some("Redo the last undone prompt edit".into()); redo.keys.push(crate::key!('r', CONTROL)); push_pseudo(&mut entries, redo, Some(REDO_LONG_HELP)); + + // Prompt history (Up / /history). Not part of the shared paste/undo/redo + // `dimmed`: that also lights on DashboardFocused, but Up-history is + // prompt-only, so give it its own PromptFocused-scoped dim. + let mut history = HintItem::new(crate::key!(Up), "history"); + history.description = Some("Prompt history".into()); + let history_dimmed = !active_contexts.contains(&When::PromptFocused); + entries.push(ShortcutsHelpEntry::Hint { + item: history, + dimmed: history_dimmed, + action_id: None, + long_help: Some(HISTORY_LONG_HELP), + }); } let count = entries.len() - header_idx - 1; if count == 0 { @@ -556,7 +604,7 @@ impl ShortcutsHelpMode { /// Build detail mode state from a cheatsheet entry (title/keys/body for the man page). /// /// Registry rows always open. Pseudo-rows (`action_id: None`) open only when they -/// ship `long_help` so list-only rows like scrollback search stay browse-only. +/// ship `long_help`; one without it stays list-only (browse-only). pub fn detail_from_entry(entry: &ShortcutsHelpEntry) -> Option { let ShortcutsHelpEntry::Hint { item, @@ -1902,6 +1950,26 @@ mod tests { }) } + fn has_find_search(entries: &[ShortcutsHelpEntry]) -> bool { + entries.iter().any(|e| { + matches!( + e, + ShortcutsHelpEntry::Hint { item, .. } + if item.custom_display == Some("/find") + ) + }) + } + + fn history_row(entries: &[ShortcutsHelpEntry]) -> Option<&ShortcutsHelpEntry> { + entries.iter().find(|e| { + matches!( + e, + ShortcutsHelpEntry::Hint { item, action_id: None, .. } + if item.label == "history" + ) + }) + } + #[test] fn build_entries_includes_scrollback_search_in_vim_mode() { let registry = ActionRegistry::defaults(); @@ -1910,15 +1978,71 @@ mod tests { has_scrollback_search(&entries), "vim cheatsheet should list / search" ); + assert!( + !has_find_search(&entries), + "vim mode uses the `/` key row, not the /find slash row" + ); } #[test] - fn build_entries_omits_scrollback_search_in_simple_mode() { + fn build_entries_includes_find_search_in_simple_mode() { let registry = ActionRegistry::defaults(); let entries = build_entries(&all_contexts(), ®istry, false); + assert!( + has_find_search(&entries), + "simple mode should list the /find scrollback search" + ); assert!( !has_scrollback_search(&entries), - "simple mode does not bind / to search, so it must not be listed" + "simple mode must not list the bare `/` key row" + ); + } + + #[test] + fn build_entries_includes_history_row_in_both_modes() { + let registry = ActionRegistry::defaults(); + for vim in [true, false] { + let entries = build_entries(&all_contexts(), ®istry, vim); + assert!( + history_row(&entries).is_some(), + "history row should appear in vim={vim} mode" + ); + } + } + + #[test] + fn history_row_lit_only_by_prompt_focus() { + let registry = ActionRegistry::defaults(); + + let entries = build_entries(&[When::PromptFocused], ®istry, false); + let ShortcutsHelpEntry::Hint { dimmed, .. } = + history_row(&entries).expect("history row present") + else { + unreachable!(); + }; + assert!( + !*dimmed, + "history row must be lit when the prompt is focused" + ); + + let entries = build_entries(&[When::ScrollbackFocused], ®istry, false); + let ShortcutsHelpEntry::Hint { dimmed, .. } = + history_row(&entries).expect("history row present") + else { + unreachable!(); + }; + assert!(*dimmed, "history row must be dimmed without prompt focus"); + + // Dashboard focus alone must not light it (unlike paste/undo/redo). + let entries = build_entries(&[When::DashboardFocused], ®istry, false); + let ShortcutsHelpEntry::Hint { dimmed, .. } = + history_row(&entries).expect("history row present") + else { + unreachable!(); + }; + assert!( + *dimmed, + "dashboard focus alone must not light the history row" ); } @@ -2170,16 +2294,27 @@ mod tests { fn build_entries_overlay_stop_wins_dedup_and_shadows_cheatsheet_ctrl_x() { let registry = ActionRegistry::defaults(); let ctrl_x = crate::key!('x', CONTROL); + // Match the two Ctrl+X rows by ActionId: the list and overlay + // stops carry different labels ("delete" vs "stop"). + let is_stop = |action_id: &Option| { + matches!( + action_id, + Some(ActionId::DashboardStop | ActionId::DashboardOverlayStop) + ) + }; let stop_rows = |entries: &[ShortcutsHelpEntry]| -> Vec<(String, bool)> { entries .iter() .filter_map(|e| match e { - ShortcutsHelpEntry::Hint { item, dimmed, .. } if item.label == "stop" => { - Some(( - item.description.as_deref().unwrap_or_default().to_string(), - *dimmed, - )) - } + ShortcutsHelpEntry::Hint { + item, + dimmed, + action_id, + .. + } if is_stop(action_id) => Some(( + item.description.as_deref().unwrap_or_default().to_string(), + *dimmed, + )), _ => None, }) .collect() @@ -2188,9 +2323,9 @@ mod tests { entries .iter() .find_map(|e| match e { - ShortcutsHelpEntry::Hint { - item, action_id, .. - } if item.label == "stop" => Some(*action_id), + ShortcutsHelpEntry::Hint { action_id, .. } if is_stop(action_id) => { + Some(*action_id) + } _ => None, }) .flatten() @@ -2212,8 +2347,7 @@ mod tests { let list = build_entries(&[When::DashboardFocused, When::Always], ®istry, true); assert_eq!( stop_rows(&list), - vec![("Stop / Close agent".to_string(), false)], - "the dashboard list must show exactly the list `stop`, lit", + vec![("Stop / Delete agent".to_string(), false)], ); assert_eq!( stop_id(&list), @@ -2729,7 +2863,7 @@ mod tests { /// Search has no long_help — Enter stays in browse. #[test] - fn enter_on_search_pseudo_row_does_not_open_detail() { + fn enter_on_search_pseudo_row_opens_detail() { let registry = ActionRegistry::defaults(); let entries = build_entries(&all_contexts(), ®istry, true); let idx = entries @@ -2737,11 +2871,24 @@ mod tests { .position(|e| { matches!( e, - ShortcutsHelpEntry::Hint { item, action_id: None, .. } - if item.label == "search" + ShortcutsHelpEntry::Hint { + item, + action_id: None, + long_help: Some(_), + .. + } if item.label == "search" ) }) .expect("vim-mode entries include the `/`-search pseudo-row"); + assert_eq!( + detail_from_entry(&entries[idx]) + .and_then(|m| match m { + ShortcutsHelpMode::Detail { body, .. } => Some(body), + _ => None, + }) + .as_deref(), + Some(SCROLLBACK_SEARCH_LONG_HELP) + ); let mut state = build_initial_picker_state(&entries); state.selected = idx; let mut mode = browse_mode(); @@ -2754,11 +2901,8 @@ mod tests { &no_expanded(), &mut mode, ); - assert_eq!(out, ShortcutsHelpOutcome::Unchanged); - assert!( - mode.is_browse(), - "search pseudo-row Enter must not open detail" - ); + assert_eq!(out, ShortcutsHelpOutcome::Changed); + assert!(mode.is_detail(), "search pseudo-row Enter opens detail"); } #[test] @@ -3273,6 +3417,9 @@ mod tests { let paste_key = key!('v', CONTROL); let undo_key = key!('z', CONTROL); let redo_key = key!('z', CONTROL | SHIFT); + // Prompt history (Up / /history) is an inline key handler + slash + // command, not an ActionRegistry entry, so it stays display-only too. + let history_key = key!(Up); for entry in &entries { let ShortcutsHelpEntry::Hint { item, action_id, .. @@ -3285,6 +3432,7 @@ mod tests { "paste" => item.keys.contains(&paste_key), "undo" => item.keys.contains(&undo_key), "redo" => item.keys.contains(&redo_key), + "history" => item.keys.contains(&history_key), _ => false, }; if is_pseudo { @@ -3407,7 +3555,7 @@ mod tests { } #[test] - fn search_pseudo_row_does_not_expand() { + fn search_pseudo_row_expands() { use crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; let registry = ActionRegistry::defaults(); let entries = build_entries(&all_contexts(), ®istry, true); @@ -3437,8 +3585,8 @@ mod tests { ); assert_eq!( out, - ShortcutsHelpOutcome::Unchanged, - "search pseudo-row must stay inert for {code:?}, got {out:?}" + ShortcutsHelpOutcome::ToggleExpand(ExpandKey::Pseudo("search")), + "search pseudo-row must expand for {code:?}, got {out:?}" ); } } diff --git a/crates/codegen/xai-grok-pager/src/views/welcome/hero_box.rs b/crates/codegen/xai-grok-pager/src/views/welcome/hero_box.rs index 527ab7c..31afeaa 100644 --- a/crates/codegen/xai-grok-pager/src/views/welcome/hero_box.rs +++ b/crates/codegen/xai-grok-pager/src/views/welcome/hero_box.rs @@ -290,6 +290,8 @@ pub(super) struct HeroBoxRects { pub(super) announcement_rect: Option, /// Promo upgrade CTA `[label]` button rect (click → open), if drawn. pub(super) upgrade_cta_rect: Option, + #[cfg(feature = "local-workspace")] + pub(super) workspace_mode_rects: super::WorkspaceModeHitRects, } /// Render the bordered hero box with logo left, version + subtitle + menu right. @@ -306,6 +308,11 @@ pub(super) fn render_hero_box( changelog_bullets: &[String], changelog_has_full_notes: bool, upgrade_cta: Option<&str>, + #[cfg(feature = "local-workspace")] workspace_mode: Option<( + super::WelcomeWorkspaceMode, + bool, + bool, + )>, ) -> HeroBoxRects { // Dim the box border toward the background for a softer, dimmer gray. let border_color = crate::render::color::blend_color(theme.bg_base, theme.gray_dim, 0.45) @@ -371,14 +378,45 @@ pub(super) fn render_hero_box( } } + #[cfg(feature = "local-workspace")] + let (menu_area, workspace_mode_rects) = + if let Some((mode, locked, ack_pending)) = workspace_mode { + let picker_rect = Rect { + height: 1.min(layout.hero_menu.height), + ..layout.hero_menu + }; + let rects = super::render_workspace_mode_picker( + picker_rect, + buf, + theme, + mode, + mouse_pos, + locked, + ack_pending, + ); + let menu_area = Rect { + y: layout.hero_menu.y + super::workspace_mode::WORKSPACE_MODE_MENU_ROWS, + height: layout + .hero_menu + .height + .saturating_sub(super::workspace_mode::WORKSPACE_MODE_MENU_ROWS), + ..layout.hero_menu + }; + (menu_area, rects) + } else { + (layout.hero_menu, super::WorkspaceModeHitRects::default()) + }; + #[cfg(not(feature = "local-workspace"))] + let menu_area = layout.hero_menu; + let menu_rects = super::menu::render_menu( - layout.hero_menu, + menu_area, buf, theme, menu_items, selected, mouse_pos, - layout.hero_menu.width, + menu_area.width, ); HeroBoxRects { menu_rects, @@ -386,6 +424,8 @@ pub(super) fn render_hero_box( announcement_truncated, announcement_rect, upgrade_cta_rect, + #[cfg(feature = "local-workspace")] + workspace_mode_rects, } } diff --git a/crates/codegen/xai-grok-pager/src/views/welcome/mod.rs b/crates/codegen/xai-grok-pager/src/views/welcome/mod.rs index fb2bd3d..2d2bc3e 100644 --- a/crates/codegen/xai-grok-pager/src/views/welcome/mod.rs +++ b/crates/codegen/xai-grok-pager/src/views/welcome/mod.rs @@ -24,6 +24,8 @@ mod menu; mod prompt; mod toast; mod top_bar; +#[cfg(feature = "local-workspace")] +pub(crate) mod workspace_mode; pub(crate) use logo::shimmer_frame; use logo::{logo_line_count, render_logo}; @@ -31,6 +33,11 @@ use menu::render_menu; pub(crate) use toast::paint_welcome_toast; pub(crate) use top_bar::location_line_at; use top_bar::render_top_bar; +#[cfg(feature = "local-workspace")] +pub use workspace_mode::{ + WelcomeWorkspaceMode, WorkspaceModeHitRects, hit_test_workspace_mode, + render_workspace_mode_picker, +}; /// True for VS Code and xterm.js embeds (VS Code-family IDEs and Zed) where /// quit is `Ctrl+D` (canonical: [`TerminalName::is_vscode_family`]). @@ -123,6 +130,9 @@ pub struct WelcomeRenderResult { pub privacy_banner_opt_out_rect: Option, pub privacy_banner_terms_rect: Option, pub privacy_banner_policy_rect: Option, + /// Hit-test rects for the chat workspace-mode segmented control. + #[cfg(feature = "local-workspace")] + pub workspace_mode_rects: WorkspaceModeHitRects, } use hero_box::HERO_BOX_MIN_WIDTH; @@ -631,6 +641,7 @@ pub struct WelcomeRenderParams<'a> { pub session_picker_grouped: bool, /// Source filter for the session picker. pub session_picker_source_filter: crate::views::session_picker::SourceFilter, + pub session_picker_pending_delete: bool, /// Process-wide `--chat`: the picker lists backend conversations only, so /// the source filter and local deep search are hidden. pub chat_mode: bool, @@ -656,6 +667,15 @@ pub struct WelcomeRenderParams<'a> { pub upgrade_cta: Option<&'a str>, /// Non-blocking welcome privacy banner above the prompt. pub privacy_banner: bool, + /// Chat-mode workspace picker selection (`local-workspace` feature). + #[cfg(feature = "local-workspace")] + pub workspace_mode: WelcomeWorkspaceMode, + /// CLI/env already stamped local workspace — picker is display-only. + #[cfg(feature = "local-workspace")] + pub workspace_mode_startup_locked: bool, + /// In-TUI ACK confirm pending for Local. + #[cfg(feature = "local-workspace")] + pub workspace_mode_ack_pending: bool, } /// Render the welcome screen. @@ -720,22 +740,7 @@ pub fn render_welcome( 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, - announcement_truncated: false, - announcement_rect: None, - upgrade_cta_rect: None, - privacy_banner_opt_in_rect: None, - privacy_banner_opt_out_rect: None, - privacy_banner_terms_rect: None, - privacy_banner_policy_rect: None, + ..Default::default() } } AuthState::Authenticating { auth_url, mode, .. } => { @@ -753,25 +758,9 @@ pub fn render_welcome( params.show_raw_url, ); WelcomeRenderResult { - cursor_pos: None, - post_flush_escapes: None, - menu_rects: vec![], - prompt_rect: None, - session_picker_hit_areas: None, - 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, - announcement_truncated: false, - announcement_rect: None, - upgrade_cta_rect: None, - privacy_banner_opt_in_rect: None, - privacy_banner_opt_out_rect: None, - privacy_banner_terms_rect: None, - privacy_banner_policy_rect: None, + ..Default::default() } } AuthState::Done if params.is_zdr_blocked => { @@ -790,25 +779,9 @@ pub fn render_welcome( 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, - announcement_truncated: false, - announcement_rect: None, - upgrade_cta_rect: None, - privacy_banner_opt_in_rect: None, - privacy_banner_opt_out_rect: None, - privacy_banner_terms_rect: None, - privacy_banner_policy_rect: None, + ..Default::default() } } // Folder-trust question: shown after auth, before any session is @@ -1792,10 +1765,25 @@ fn render_welcome_done( owned_menu.as_slice() }; + #[cfg(feature = "local-workspace")] + // Keep the segmented control (and ACK y/N) visible when history is open + // if first-run Local ACK is pending — otherwise the confirm is unpainted + // while the ACK handler still swallows keys. + let show_workspace_picker = + p.chat_mode && p.has_access && (!show_picker || p.workspace_mode_ack_pending); + #[cfg(feature = "local-workspace")] + let workspace_picker_rows = if show_workspace_picker { + workspace_mode::WORKSPACE_MODE_MENU_ROWS + } else { + 0 + }; + #[cfg(not(feature = "local-workspace"))] + let workspace_picker_rows = 0u16; + let menu_height = if show_picker { 0 } else { - menu_items.len() as u16 + menu_items.len() as u16 + workspace_picker_rows }; // Session picker height: 1 row per entry (no dividers), scrollable. @@ -1843,6 +1831,8 @@ fn render_welcome_done( let mut announcement_rect: Option = None; let mut upgrade_cta_rect: Option = None; + #[cfg(feature = "local-workspace")] + let mut workspace_mode_rects = WorkspaceModeHitRects::default(); let (menu_rects, picker_close_button) = if show_picker { // Use the full area since logo/menu are hidden and shortcuts // are now rendered inside the picker content area. @@ -1868,6 +1858,7 @@ fn render_welcome_done( tick: p.welcome_tick, grouped: p.session_picker_grouped, source_filter: p.session_picker_source_filter, + pending_delete: p.session_picker_pending_delete, chat_mode: p.chat_mode, cwd: p.cwd, }, @@ -1887,11 +1878,21 @@ fn render_welcome_done( p.changelog_bullets, p.changelog_has_full_notes, p.upgrade_cta, + #[cfg(feature = "local-workspace")] + show_workspace_picker.then_some(( + p.workspace_mode, + p.workspace_mode_startup_locked, + p.workspace_mode_ack_pending, + )), ); changelog_cta_rect = rects.changelog_cta_rect; announcement_truncated = rects.announcement_truncated; announcement_rect = rects.announcement_rect; upgrade_cta_rect = rects.upgrade_cta_rect; + #[cfg(feature = "local-workspace")] + { + workspace_mode_rects = rects.workspace_mode_rects; + } (rects.menu_rects, None) } else { // Narrow layout: stacked logo above, menu below. Inset the menu the @@ -1899,6 +1900,28 @@ fn render_welcome_done( // instead of touching the window edge on narrow terminals. render_logo(layout.logo, buf, theme, content_area.height); let menu_area = inset_horizontal(layout.menu, prompt::prompt_inset(p.compact)); + #[cfg(feature = "local-workspace")] + let menu_area = if show_workspace_picker { + let picker_rect = workspace_mode::picker_area(menu_area); + workspace_mode_rects = render_workspace_mode_picker( + picker_rect, + buf, + theme, + p.workspace_mode, + p.mouse_pos, + p.workspace_mode_startup_locked, + p.workspace_mode_ack_pending, + ); + Rect { + y: menu_area.y + workspace_mode::WORKSPACE_MODE_MENU_ROWS, + height: menu_area + .height + .saturating_sub(workspace_mode::WORKSPACE_MODE_MENU_ROWS), + ..menu_area + } + } else { + menu_area + }; ( render_menu( menu_area, @@ -2228,6 +2251,8 @@ fn render_welcome_done( privacy_banner_opt_out_rect, privacy_banner_terms_rect, privacy_banner_policy_rect, + #[cfg(feature = "local-workspace")] + workspace_mode_rects, } } @@ -2252,6 +2277,7 @@ pub(crate) struct SessionPickerRenderCtx<'a> { pub(crate) grouped: bool, /// Source filter for filtering session entries. pub(crate) source_filter: crate::views::session_picker::SourceFilter, + pub(crate) pending_delete: bool, /// Process-wide `--chat`: hides the source-filter chip and the /// deep-search/filter footer hints (see `WelcomeRenderParams::chat_mode`). pub(crate) chat_mode: bool, @@ -2448,7 +2474,23 @@ pub(crate) fn render_session_picker( description: None, pinned: false, }); - if !ctx.chat_mode { + if ctx.pending_delete { + default_shortcuts.clear(); + default_shortcuts.push(HintItem { + keys: vec![], + label: "confirm delete".into(), + custom_display: Some("y"), + description: None, + pinned: false, + }); + default_shortcuts.push(HintItem { + keys: vec![], + label: "cancel".into(), + custom_display: Some("n"), + description: None, + pinned: false, + }); + } else if !ctx.chat_mode { default_shortcuts.push(HintItem { keys: vec![], label: "filter".into(), @@ -2456,6 +2498,13 @@ pub(crate) fn render_session_picker( description: None, pinned: false, }); + default_shortcuts.push(HintItem { + keys: vec![], + label: "delete".into(), + custom_display: Some("d"), + description: None, + pinned: false, + }); } let config = PickerConfig { @@ -2474,7 +2523,11 @@ pub(crate) fn render_session_picker( filter_key_hint: (!ctx.chat_mode).then_some("f"), filter_active: !ctx.chat_mode && ctx.source_filter.is_active(), header_note: hidden_hint.as_deref(), - action_keys: &[], + action_keys: if ctx.chat_mode || ctx.pending_delete { + &[] + } else { + &[('d', "delete")] + }, disable_search: false, compact_bottom_bar: false, search_only_on_slash: false, @@ -2789,6 +2842,7 @@ mod tests { subscription_tier: None, session_picker_grouped: false, session_picker_source_filter: crate::views::session_picker::SourceFilter::default(), + session_picker_pending_delete: false, chat_mode: false, cwd: std::path::Path::new("/repo"), credit_balance: None, @@ -2799,6 +2853,12 @@ mod tests { welcome_announcement_expanded: false, upgrade_cta: None, privacy_banner: false, + #[cfg(feature = "local-workspace")] + workspace_mode: WelcomeWorkspaceMode::Sandbox, + #[cfg(feature = "local-workspace")] + workspace_mode_startup_locked: false, + #[cfg(feature = "local-workspace")] + workspace_mode_ack_pending: false, } } @@ -2964,6 +3024,7 @@ mod tests { tick: 0, grouped: false, source_filter: crate::views::session_picker::SourceFilter::default(), + pending_delete: false, chat_mode: true, }, ); @@ -3039,6 +3100,7 @@ mod tests { tick: 0, grouped: false, source_filter: crate::views::session_picker::SourceFilter::default(), + pending_delete: false, chat_mode, }, ); diff --git a/crates/codegen/xai-grok-pager/src/views/welcome/workspace_mode.rs b/crates/codegen/xai-grok-pager/src/views/welcome/workspace_mode.rs new file mode 100644 index 0000000..2f1b2bf --- /dev/null +++ b/crates/codegen/xai-grok-pager/src/views/welcome/workspace_mode.rs @@ -0,0 +1,841 @@ +//! Welcome Sandbox | Local picker under `--chat`. CLI/env stamp wins at startup. + +use ratatui::buffer::Buffer; +use ratatui::layout::{Constraint, Flex, Layout, Position, Rect}; +use ratatui::style::{Modifier, Style}; +use ratatui::text::Span; +use unicode_width::UnicodeWidthStr; + +use crate::theme::Theme; + +/// Welcome-screen workspace selection (in-memory until session start). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum WelcomeWorkspaceMode { + /// Backend sandbox / product-chat default. + #[default] + Sandbox, + /// Local Computer Hub workspace server (own mode; replaces sandbox). + LocalWorkspace, +} + +impl WelcomeWorkspaceMode { + /// Modes shown on the welcome picker under `--chat`. + pub const ALL: [Self; 2] = [Self::Sandbox, Self::LocalWorkspace]; + + pub fn cycle_next(self) -> Self { + match self { + Self::Sandbox => Self::LocalWorkspace, + Self::LocalWorkspace => Self::Sandbox, + } + } + + pub fn cycle_prev(self) -> Self { + self.cycle_next() + } + + pub fn label(self) -> &'static str { + match self { + Self::Sandbox => "Sandbox", + Self::LocalWorkspace => "Local workspace", + } + } + + pub fn hint(self) -> &'static str { + match self { + Self::Sandbox => "backend sandbox", + Self::LocalWorkspace => "this machine · Computer Hub", + } + } + + /// Compact in-session / status-bar label. + pub fn status_label(self, cli_locked: bool) -> &'static str { + match (self, cli_locked) { + (Self::Sandbox, _) => "Sandbox", + (Self::LocalWorkspace, true) => "Local·CLI", + (Self::LocalWorkspace, false) => "Local", + } + } + + /// Unified-list `kind`: Sandbox → `chat`, Local → `build`. + pub fn history_kind_filter(self) -> &'static str { + match self { + Self::Sandbox => "chat", + Self::LocalWorkspace => "build", + } + } + + /// Conversation/gateway → Sandbox; other sources → Local. + pub fn from_history_source(source: &str) -> Self { + if source == "conversation" { + Self::Sandbox + } else { + Self::LocalWorkspace + } + } + + pub fn index(self) -> usize { + match self { + Self::Sandbox => 0, + Self::LocalWorkspace => 1, + } + } + + pub fn from_index(i: usize) -> Self { + Self::ALL[i % Self::ALL.len()] + } +} + +/// Structured log target for welcome / in-session workspace mode events. +pub const WORKSPACE_MODE_LOG: &str = "grok.pager.workspace_mode"; + +/// Log a welcome picker selection change (Ctrl+E cycle or click). +pub fn log_welcome_mode_selected( + mode: WelcomeWorkspaceMode, + via: &'static str, + startup_locked: bool, +) { + tracing::info!( + target: WORKSPACE_MODE_LOG, + event = "welcome_mode_selected", + mode = mode.label(), + history_kind = mode.history_kind_filter(), + via, + startup_locked, + "welcome workspace mode selected" + ); +} + +/// Log Local ACK confirm or cancel. +pub fn log_welcome_ack(outcome: &'static str) { + tracing::info!( + target: WORKSPACE_MODE_LOG, + event = "welcome_local_ack", + outcome, + "welcome local-workspace ACK" + ); +} + +/// Log one-shot / process stamp application for a new welcome session. +pub fn log_welcome_intent_applied( + mode: WelcomeWorkspaceMode, + startup_locked: bool, + one_shot: &'static str, + process_stamp: &'static str, +) { + tracing::info!( + target: WORKSPACE_MODE_LOG, + event = "welcome_intent_applied", + mode = mode.label(), + startup_locked, + one_shot, + process_stamp, + "welcome workspace intent applied for NewSession" + ); +} + +/// Log CLI/env lock applied at startup (before any welcome selection). +pub fn log_cli_lock_applied(mode: WelcomeWorkspaceMode) { + tracing::info!( + target: WORKSPACE_MODE_LOG, + event = "cli_lock_applied", + mode = mode.label(), + "CLI/env local-workspace lock applied at startup" + ); +} + +/// Log CLI/env lock winning over a differing welcome selection. +pub fn log_cli_lock_wins(mode: WelcomeWorkspaceMode) { + tracing::info!( + target: WORKSPACE_MODE_LOG, + event = "cli_lock_wins", + mode = mode.label(), + "CLI/env local-workspace lock wins; welcome selection ignored" + ); +} + +/// In-session indicator: history bypass / local intent → Local; else Sandbox. +pub fn indicator_for_opening_session( + chat_kind: bool, + history_load_as_build: bool, + cli_locked: bool, + local_workspace_intent: bool, +) -> (WelcomeWorkspaceMode, bool) { + if history_load_as_build { + return (WelcomeWorkspaceMode::LocalWorkspace, cli_locked); + } + if chat_kind && !local_workspace_intent { + return (WelcomeWorkspaceMode::Sandbox, false); + } + if cli_locked { + return (WelcomeWorkspaceMode::LocalWorkspace, true); + } + if local_workspace_intent { + return (WelcomeWorkspaceMode::LocalWorkspace, false); + } + (WelcomeWorkspaceMode::Sandbox, false) +} + +/// Log session-list kind filter / history-source switch. +pub fn log_history_source( + event: &'static str, + mode: Option, + kind_filter: Option<&[String]>, + source: Option<&str>, +) { + tracing::info!( + target: WORKSPACE_MODE_LOG, + event, + mode = mode.map(WelcomeWorkspaceMode::label), + kind_filter = ?kind_filter, + history_source = source, + "workspace history source" + ); +} + +/// Hit-test rects for each segmented option (Sandbox, Local). +#[derive(Debug, Clone, Default)] +pub struct WorkspaceModeHitRects { + pub options: [Option; 2], + pub row: Option, +} + +/// Rows reserved above the welcome menu for the picker (content + gap). +pub const WORKSPACE_MODE_MENU_ROWS: u16 = 2; + +/// Paint the segmented workspace control into `area`. +/// +/// Layout: +/// `Workspace [ Sandbox ] [ Local workspace ] ctrl+e` +/// or when locked: `Workspace [ Local workspace ] locked by CLI` +pub fn render_workspace_mode_picker( + area: Rect, + buf: &mut Buffer, + theme: &Theme, + selected: WelcomeWorkspaceMode, + mouse_pos: Option<(u16, u16)>, + startup_locked: bool, + ack_pending: bool, +) -> WorkspaceModeHitRects { + if area.height == 0 || area.width < 20 { + return WorkspaceModeHitRects::default(); + } + + let row = Rect { + x: area.x, + y: area.y, + width: area.width, + height: 1, + }; + + let label_style = Style::default().fg(theme.gray); + let key_style = Style::default().fg(theme.gray_bright); + let inactive = Style::default().fg(theme.gray_bright); + let active = Style::default() + .fg(theme.bg_base) + .bg(theme.accent_user) + .add_modifier(Modifier::BOLD); + let hover = Style::default() + .fg(theme.text_primary) + .add_modifier(Modifier::BOLD); + let locked_style = Style::default().fg(theme.gray); + + buf.set_span(row.x, row.y, &Span::styled("Workspace ", label_style), 11); + + let mut x = row.x.saturating_add(11); + let mut options = [None; 2]; + + let modes: &[WelcomeWorkspaceMode] = if startup_locked { + // Locked: show the effective mode only (CLI/env stamp). + match selected { + WelcomeWorkspaceMode::LocalWorkspace => &[WelcomeWorkspaceMode::LocalWorkspace], + WelcomeWorkspaceMode::Sandbox => &[WelcomeWorkspaceMode::Sandbox], + } + } else { + &WelcomeWorkspaceMode::ALL + }; + + for (slot, mode) in modes.iter().enumerate() { + if x >= row.x + row.width { + break; + } + let text = if *mode == selected { + format!(" • {} ", mode.label()) + } else { + format!(" {} ", mode.label()) + }; + let w = UnicodeWidthStr::width(text.as_str()) as u16; + if x + w > row.x + row.width { + break; + } + let rect = Rect { + x, + y: row.y, + width: w, + height: 1, + }; + let hovered = !startup_locked + && !ack_pending + && mouse_pos.is_some_and(|(mx, my)| rect.contains(Position::new(mx, my))); + let style = if *mode == selected { + active + } else if hovered { + hover + } else { + inactive + }; + buf.set_span(x, row.y, &Span::styled(text, style), w); + if slot < options.len() { + // Map by mode index so hit-test stays stable. + options[mode.index()] = Some(rect); + } + x = x.saturating_add(w); + if slot + 1 < modes.len() && x + 1 < row.x + row.width { + buf.set_span(x, row.y, &Span::styled("│", label_style), 1); + x = x.saturating_add(1); + } + } + + let trailing = if ack_pending { + " confirm local workspace? y/N" + } else if startup_locked { + " locked by CLI" + } else { + " ctrl+e" + }; + let trailing_style = if ack_pending { + Style::default() + .fg(theme.text_primary) + .add_modifier(Modifier::BOLD) + } else if startup_locked { + locked_style + } else { + key_style + }; + if !trailing.is_empty() && x + trailing.len() as u16 <= row.x + row.width { + buf.set_span( + row.x + row.width - trailing.len() as u16, + row.y, + &Span::styled(trailing, trailing_style), + trailing.len() as u16, + ); + } else if ack_pending && row.width > 20 { + // Narrow terminals: paint confirm over the right side so it stays visible. + let short = " y/N confirm local"; + let start = row.x + row.width.saturating_sub(short.len() as u16); + buf.set_span( + start, + row.y, + &Span::styled(short, trailing_style), + short.len() as u16, + ); + } + + WorkspaceModeHitRects { + options, + row: Some(row), + } +} + +/// Hit-test a click against option rects. Returns the selected mode if hit. +pub fn hit_test_workspace_mode( + rects: &WorkspaceModeHitRects, + column: u16, + row: u16, +) -> Option { + let pos = Position::new(column, row); + for (i, rect) in rects.options.iter().enumerate() { + if rect.is_some_and(|r| r.contains(pos)) { + return Some(WelcomeWorkspaceMode::from_index(i)); + } + } + None +} + +/// Result of preparing welcome workspace intent for a new session. +#[cfg(feature = "local-workspace")] +#[derive(Debug)] +pub enum WelcomeWorkspacePrepare { + /// Continue. `session_override`: `Some(None)` sandbox, `Some(Some)` local, `None` keep stamp. + Continue { + session_override: Option>, + warning: Option, + }, + /// Stay on welcome; show in-TUI ACK confirm before stamping Local. + AwaitAck, +} + +/// Prepare welcome Sandbox/Local for NewSession. Local may return `AwaitAck`. +#[cfg(feature = "local-workspace")] +pub fn prepare_welcome_workspace_for_new_session( + selection: WelcomeWorkspaceMode, + startup_locked: bool, + chat_mode: bool, + cwd: &std::path::Path, + agents_alive: bool, +) -> anyhow::Result { + use crate::app::session_startup::{ + local_workspace_ack_satisfied, resolve_local_workspace_config, set_active_local_workspace, + }; + + if startup_locked || !chat_mode { + if startup_locked { + log_cli_lock_wins(selection); + } + return Ok(WelcomeWorkspacePrepare::Continue { + session_override: None, + warning: None, + }); + } + + match selection { + WelcomeWorkspaceMode::Sandbox => { + if !agents_alive { + // Safe: no live session still reading the process stamp. + set_active_local_workspace(None)?; + } + log_welcome_intent_applied( + selection, + startup_locked, + "sandbox_none", + if agents_alive { "kept" } else { "cleared" }, + ); + Ok(WelcomeWorkspacePrepare::Continue { + session_override: Some(None), + warning: None, + }) + } + WelcomeWorkspaceMode::LocalWorkspace => { + if !local_workspace_ack_satisfied() { + tracing::info!( + target: WORKSPACE_MODE_LOG, + event = "welcome_local_ack", + outcome = "await", + "welcome Local requires ACK confirm" + ); + return Ok(WelcomeWorkspacePrepare::AwaitAck); + } + let cfg = resolve_local_workspace_config(true, Some(None), None, Some(cwd))? + .ok_or_else(|| { + anyhow::anyhow!( + "local-workspace resolve returned no config after own-mode request" + ) + })?; + // Live sessions still read the process stamp; oneshot-only then. + if !agents_alive { + set_active_local_workspace(Some(cfg.clone()))?; + } + log_welcome_intent_applied( + selection, + startup_locked, + "own_oneshot", + if agents_alive { "kept" } else { "stamped_own" }, + ); + Ok(WelcomeWorkspacePrepare::Continue { + session_override: Some(Some(cfg)), + warning: None, + }) + } + } +} + +/// Confirm Local ACK. If `agents_alive`, return oneshot only (keep process stamp). +#[cfg(feature = "local-workspace")] +pub fn confirm_welcome_local_workspace_ack( + cwd: &std::path::Path, + agents_alive: bool, +) -> anyhow::Result { + use crate::app::session_startup::{ + resolve_local_workspace_config, set_active_local_workspace, write_local_workspace_ack, + }; + + let cfg = resolve_local_workspace_config(true, Some(None), None, Some(cwd))? + .ok_or_else(|| anyhow::anyhow!("local-workspace resolve returned no config after ack"))?; + if !agents_alive { + set_active_local_workspace(Some(cfg.clone()))?; + } + write_local_workspace_ack(); + log_welcome_ack("confirmed"); + Ok(cfg) +} + +/// Sync UI selection from a startup-locked stamp (Own/Attach → Local). +#[cfg(feature = "local-workspace")] +pub fn mode_from_active_stamp( + stamp: Option<&crate::app::session_startup::LocalWorkspaceConfig>, +) -> WelcomeWorkspaceMode { + match stamp { + Some(_) => WelcomeWorkspaceMode::LocalWorkspace, + None => WelcomeWorkspaceMode::Sandbox, + } +} + +/// Whether keyboard/mouse should mutate the welcome selection. +/// +/// Same surface as ACK + render: chat mode, access, auth Done, not ZDR, +/// not CLI-startup-locked, and history picker closed (Ctrl+E/click would +/// otherwise mutate with no on-screen control). +pub fn picker_interactive( + chat_mode: bool, + has_access: bool, + auth_done: bool, + zdr_blocked: bool, + session_picker_open: bool, + startup_locked: bool, +) -> bool { + chat_mode && has_access && auth_done && !zdr_blocked && !startup_locked && !session_picker_open +} + +/// Center the picker within `menu_area` the same way the menu is inset. +pub fn picker_area(menu_area: Rect) -> Rect { + let [_, centered, _] = Layout::horizontal([ + Constraint::Min(0), + Constraint::Length(menu_area.width), + Constraint::Min(0), + ]) + .flex(Flex::Start) + .areas(menu_area); + Rect { + height: 1.min(centered.height), + ..centered + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cycle_walks_sandbox_and_local() { + let mut mode = WelcomeWorkspaceMode::Sandbox; + mode = mode.cycle_next(); + assert_eq!(mode, WelcomeWorkspaceMode::LocalWorkspace); + mode = mode.cycle_next(); + assert_eq!(mode, WelcomeWorkspaceMode::Sandbox); + assert_eq!( + WelcomeWorkspaceMode::LocalWorkspace.cycle_prev(), + WelcomeWorkspaceMode::Sandbox + ); + } + + #[test] + fn labels_are_stable() { + assert_eq!(WelcomeWorkspaceMode::Sandbox.label(), "Sandbox"); + assert_eq!( + WelcomeWorkspaceMode::LocalWorkspace.label(), + "Local workspace" + ); + assert!( + WelcomeWorkspaceMode::LocalWorkspace + .hint() + .contains("Computer Hub") + ); + assert_eq!(WelcomeWorkspaceMode::Sandbox.status_label(false), "Sandbox"); + assert_eq!( + WelcomeWorkspaceMode::LocalWorkspace.status_label(false), + "Local" + ); + assert_eq!( + WelcomeWorkspaceMode::LocalWorkspace.status_label(true), + "Local·CLI" + ); + assert_eq!(WelcomeWorkspaceMode::Sandbox.history_kind_filter(), "chat"); + assert_eq!( + WelcomeWorkspaceMode::LocalWorkspace.history_kind_filter(), + "build" + ); + assert_eq!( + WelcomeWorkspaceMode::from_history_source("conversation"), + WelcomeWorkspaceMode::Sandbox + ); + assert_eq!( + WelcomeWorkspaceMode::from_history_source("local"), + WelcomeWorkspaceMode::LocalWorkspace + ); + } + + #[test] + fn index_roundtrip() { + for mode in WelcomeWorkspaceMode::ALL { + assert_eq!(WelcomeWorkspaceMode::from_index(mode.index()), mode); + } + } + + #[test] + fn hit_test_prefers_option_rects() { + let rects = WorkspaceModeHitRects { + options: [Some(Rect::new(10, 5, 9, 1)), Some(Rect::new(20, 5, 17, 1))], + row: Some(Rect::new(0, 5, 80, 1)), + }; + assert_eq!( + hit_test_workspace_mode(&rects, 12, 5), + Some(WelcomeWorkspaceMode::Sandbox) + ); + assert_eq!( + hit_test_workspace_mode(&rects, 25, 5), + Some(WelcomeWorkspaceMode::LocalWorkspace) + ); + assert_eq!(hit_test_workspace_mode(&rects, 0, 5), None); + assert_eq!(hit_test_workspace_mode(&rects, 12, 6), None); + } + + #[test] + fn render_assigns_option_rects() { + let area = Rect::new(0, 0, 100, 2); + let mut buf = Buffer::empty(area); + let theme = Theme::current(); + let hits = render_workspace_mode_picker( + area, + &mut buf, + &theme, + WelcomeWorkspaceMode::LocalWorkspace, + None, + false, + false, + ); + assert!(hits.options[0].is_some()); + assert!(hits.options[1].is_some()); + assert!(hits.row.is_some()); + let cell = buf.cell((0, 0)).expect("cell"); + assert_eq!(cell.symbol(), "W"); + let selected = hits.options[1].expect("local selected rect"); + let selected_text = format!(" • {} ", WelcomeWorkspaceMode::LocalWorkspace.label()); + assert_eq!( + selected.width, + UnicodeWidthStr::width(selected_text.as_str()) as u16, + "option width must be display columns, not UTF-8 bytes" + ); + assert!( + selected.width < selected_text.len() as u16, + "bullet U+2022 is 3 bytes / 1 column: {selected_text:?}" + ); + } + + #[test] + fn render_ack_pending_shows_durable_confirm() { + let area = Rect::new(0, 0, 120, 1); + let mut buf = Buffer::empty(area); + let theme = Theme::current(); + let _ = render_workspace_mode_picker( + area, + &mut buf, + &theme, + WelcomeWorkspaceMode::LocalWorkspace, + None, + false, + true, + ); + let line: String = (0..area.width) + .filter_map(|x| buf.cell((x, 0)).map(|c| c.symbol().to_string())) + .collect(); + assert!( + line.contains("y/N") || line.contains("confirm"), + "ack-pending UI must stay visible: {line:?}" + ); + } + + #[test] + fn picker_interactive_matrix() { + assert!(picker_interactive(true, true, true, false, false, false)); + assert!(!picker_interactive(true, true, true, false, false, true)); + assert!( + !picker_interactive(true, true, true, false, true, false), + "history open: Ctrl+E/click must not mutate a hidden control" + ); + assert!(!picker_interactive(true, false, true, false, false, false)); + assert!(!picker_interactive(false, true, true, false, false, false)); + assert!( + !picker_interactive(true, true, false, false, false, false), + "login / authenticating must not cycle mode" + ); + assert!( + !picker_interactive(true, true, true, true, false, false), + "ZDR-blocked welcome must not cycle mode" + ); + } + + #[test] + fn indicator_derives_from_opened_session() { + assert_eq!( + indicator_for_opening_session(true, false, false, false), + (WelcomeWorkspaceMode::Sandbox, false) + ); + assert_eq!( + indicator_for_opening_session(false, true, false, false), + (WelcomeWorkspaceMode::LocalWorkspace, false) + ); + // Conversation / chat_kind without this-session local intent → Sandbox + // even when the process has a CLI lock (LoadSession strips stamp). + assert_eq!( + indicator_for_opening_session(true, false, true, false), + (WelcomeWorkspaceMode::Sandbox, false) + ); + assert_eq!( + indicator_for_opening_session(true, false, true, true), + (WelcomeWorkspaceMode::LocalWorkspace, true) + ); + assert_eq!( + indicator_for_opening_session(false, true, true, false), + (WelcomeWorkspaceMode::LocalWorkspace, true) + ); + assert_eq!(WelcomeWorkspaceMode::Sandbox.status_label(true), "Sandbox"); + } +} + +#[cfg(all(test, feature = "local-workspace"))] +mod apply_tests { + use super::*; + use crate::app::session_startup::{ + GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV, LocalWorkspaceMode, set_active_local_workspace, + }; + + #[test] + fn startup_lock_skips_override() { + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + set_active_local_workspace(Some(crate::app::session_startup::LocalWorkspaceConfig { + mode: LocalWorkspaceMode::Attach, + cwd: Some(tmp.path().to_path_buf()), + server_id: Some("cli-srv".into()), + })) + .unwrap(); + + let out = prepare_welcome_workspace_for_new_session( + WelcomeWorkspaceMode::Sandbox, + true, + true, + tmp.path(), + false, + ) + .unwrap(); + match out { + WelcomeWorkspacePrepare::Continue { + session_override, .. + } => { + assert!(session_override.is_none()); + } + WelcomeWorkspacePrepare::AwaitAck => panic!("locked must continue"), + } + let stamp = crate::app::session_startup::active_local_workspace() + .unwrap() + .expect("cli stamp kept"); + assert_eq!(stamp.mode, LocalWorkspaceMode::Attach); + set_active_local_workspace(None).unwrap(); + } + + #[test] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ACK)] + fn welcome_local_one_shot_only_when_agents_alive() { + let _ack = xai_grok_test_support::EnvGuard::set(GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV, "1"); + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let out = prepare_welcome_workspace_for_new_session( + WelcomeWorkspaceMode::LocalWorkspace, + false, + true, + tmp.path(), + true, // agents alive + ) + .unwrap(); + let WelcomeWorkspacePrepare::Continue { + session_override, .. + } = out + else { + panic!("expected continue"); + }; + assert!(session_override.flatten().is_some()); + assert!( + crate::app::session_startup::active_local_workspace() + .unwrap() + .is_none(), + "must not overwrite process stamp while other agents are alive" + ); + } + + #[test] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ACK)] + fn welcome_local_stamps_own_mode() { + let _ack = xai_grok_test_support::EnvGuard::set(GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV, "1"); + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let out = prepare_welcome_workspace_for_new_session( + WelcomeWorkspaceMode::LocalWorkspace, + false, + true, + tmp.path(), + false, + ) + .unwrap(); + let WelcomeWorkspacePrepare::Continue { + session_override, .. + } = out + else { + panic!("expected continue"); + }; + let cfg = session_override.flatten().expect("own stamp override"); + assert_eq!(cfg.mode, LocalWorkspaceMode::Own); + assert_eq!(cfg.cwd.as_deref(), Some(tmp.path())); + assert!(cfg.server_id.is_none()); + set_active_local_workspace(None).unwrap(); + } + + #[test] + fn sandbox_does_not_clear_stamp_when_agents_alive() { + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + set_active_local_workspace(Some(crate::app::session_startup::LocalWorkspaceConfig { + mode: LocalWorkspaceMode::Own, + cwd: Some(tmp.path().to_path_buf()), + server_id: None, + })) + .unwrap(); + + let out = prepare_welcome_workspace_for_new_session( + WelcomeWorkspaceMode::Sandbox, + false, + true, + tmp.path(), + true, // agents alive + ) + .unwrap(); + let WelcomeWorkspacePrepare::Continue { + session_override, .. + } = out + else { + panic!("expected continue"); + }; + assert_eq!(session_override, Some(None)); + assert!( + crate::app::session_startup::active_local_workspace() + .unwrap() + .is_some(), + "process stamp must remain for live agents" + ); + set_active_local_workspace(None).unwrap(); + } + + #[test] + #[serial_test::serial(GROK_CHAT_LOCAL_WORKSPACE_ACK)] + fn local_without_ack_awaits_confirm() { + let _ack = xai_grok_test_support::EnvGuard::unset(GROK_CHAT_LOCAL_WORKSPACE_ACK_ENV); + // Isolate ack file from developer machine. + let home = tempfile::tempdir().unwrap(); + let _home = + xai_grok_test_support::EnvGuard::set("GROK_HOME", home.path().to_str().unwrap()); + set_active_local_workspace(None).unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let out = prepare_welcome_workspace_for_new_session( + WelcomeWorkspaceMode::LocalWorkspace, + false, + true, + tmp.path(), + false, + ) + .unwrap(); + assert!(matches!(out, WelcomeWorkspacePrepare::AwaitAck)); + assert!( + crate::app::session_startup::active_local_workspace() + .unwrap() + .is_none() + ); + } +} diff --git a/crates/codegen/xai-grok-sampler/src/actor/request_task.rs b/crates/codegen/xai-grok-sampler/src/actor/request_task.rs index aa829f8..d29b982 100644 --- a/crates/codegen/xai-grok-sampler/src/actor/request_task.rs +++ b/crates/codegen/xai-grok-sampler/src/actor/request_task.rs @@ -18,7 +18,7 @@ use tokio_util::sync::CancellationToken; use tracing::Instrument; use xai_grok_sampling_types::{ - ConversationRequest, ConversationResponse, EmptyResponseContext, SamplingError, + ConversationRequest, ConversationResponse, EmptyResponseContext, SamplingError, SentCredential, error::Result as SamplingResult, }; @@ -692,7 +692,10 @@ fn synthesize_from_info(info: &SamplingErrorInfo) -> SamplingError { .find_map(|tok| tok.strip_suffix('s').and_then(|n| n.parse::().ok())) .unwrap_or(0), }, - SamplingErrorKind::Auth => SamplingError::Auth(info.message.clone()), + SamplingErrorKind::Auth => SamplingError::Auth { + message: info.message.clone(), + credential: info.credential, + }, // Must stay Serialization: EventStreamError is retryable, and a // response-parse failure is deterministic on retry. `info.message` // is the variant's rendered Display, so rebuild via the constructor @@ -826,6 +829,7 @@ fn handle_cancellation( empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: SentCredential::Unknown, }; let _ = event_tx.send(SamplingEvent::Failed { request_id: request_id.clone(), @@ -833,7 +837,7 @@ fn handle_cancellation( }); send_completion( completion_tx, - Err(SamplingError::Auth("request cancelled".to_string())), + Err(SamplingError::auth_unknown("request cancelled")), ); } @@ -863,6 +867,7 @@ mod tests { empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: SentCredential::Unknown, }; let err = synthesize_from_info(&info); match err { @@ -883,6 +888,7 @@ mod tests { empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: SentCredential::Unknown, }; let err = synthesize_from_info(&info); match err { @@ -908,6 +914,7 @@ mod tests { empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: SentCredential::Unknown, }; let err = synthesize_from_info(&info); match err { diff --git a/crates/codegen/xai-grok-sampler/src/client.rs b/crates/codegen/xai-grok-sampler/src/client.rs index ec9e584..65cc6cf 100644 --- a/crates/codegen/xai-grok-sampler/src/client.rs +++ b/crates/codegen/xai-grok-sampler/src/client.rs @@ -25,8 +25,8 @@ use xai_grok_sampling_types::error::{try_parse_stream_error, user_facing_api_err use xai_grok_sampling_types::{ ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse, ConversationRequest, ConversationResponse, CreateResponseWrapper, DOOM_LOOP_CHECK_HEADER, MessagesRequestWrapper, - ResponseModelMetadata, Result, SamplingError, build_messages_request, is_check_event, messages, - rs, + ResponseModelMetadata, Result, SamplingError, SentCredential, build_messages_request, + is_check_event, messages, rs, }; use crate::attribution::bearer_tail_fragment; @@ -488,6 +488,28 @@ pub fn user_agent_string_for(origin: &OriginClientInfo) -> String { } } +/// A request builder coupled to the credential state it was built with, so +/// a 401 arm cannot classify from anything but the build-time capture. The +/// wire default (`SentCredential::Unknown`, which charges the retry budget) +/// stays the fail-closed one; only an explicit `sent_bearer: None` — a send +/// the builder provably stamped no credential onto — reaches the uncharged +/// lane via [`auth_rejected`]. +struct SentRequest { + builder: reqwest::RequestBuilder, + /// Tail fragment of the credential in the built headers (`None` = no + /// credential header at all). + sent_bearer: Option, +} + +/// The one way a 401 becomes a `SamplingError::Auth` with a wire-derived +/// credential classification: from the fragment its [`SentRequest`] captured. +fn auth_rejected(message: String, sent_bearer: Option<&str>) -> SamplingError { + SamplingError::Auth { + message, + credential: SentCredential::from_sent_fragment(sent_bearer), + } +} + // ============================================================================= // SamplingClient // ============================================================================= @@ -510,9 +532,8 @@ impl SamplingClient { api_key = %api_key, "Invalid api_key: cannot be converted to a valid HTTP header" ); - SamplingError::Auth( - "Invalid api_key: cannot be converted to a valid HTTP header" - .to_string(), + SamplingError::auth_unknown( + "Invalid api_key: cannot be converted to a valid HTTP header", ) })?; headers.insert(HeaderName::from_static("x-api-key"), header_value); @@ -524,9 +545,8 @@ impl SamplingClient { api_key = %api_key, "Invalid api_key: cannot be converted to a valid HTTP Authorization header" ); - SamplingError::Auth( - "Invalid api_key: cannot be converted to a valid HTTP Authorization header" - .to_string(), + SamplingError::auth_unknown( + "Invalid api_key: cannot be converted to a valid HTTP Authorization header", ) })?; headers.insert(AUTHORIZATION, header_value); @@ -658,7 +678,7 @@ impl SamplingClient { self.defaults.api_backend.clone() } - /// POST with default headers, returning the builder plus the tail + /// POST with default headers, returning the builder coupled to the tail /// fragment of the credential actually placed in its headers (`None` = /// no credential) — captured at build time because a record-time /// re-read races with the recovery a 401 triggers. @@ -666,7 +686,7 @@ impl SamplingClient { /// A wired bearer_resolver is the sole auth source: a missing live /// bearer strips default Authorization / x-api-key so a hard-expired /// seed key cannot ride on the wire. - fn post(&self, url: impl reqwest::IntoUrl) -> (reqwest::RequestBuilder, Option) { + fn post(&self, url: impl reqwest::IntoUrl) -> SentRequest { let mut headers = self.default_headers.clone(); if let Some(resolver) = &self.bearer_resolver { headers.remove(AUTHORIZATION); @@ -713,7 +733,10 @@ impl SamplingClient { if let Some(injector) = &self.header_injector { injector.inject(&mut headers); } - (self.http.post(url).headers(headers), sent_bearer) + SentRequest { + builder: self.http.post(url).headers(headers), + sent_bearer, + } } /// Tail fragment of the credential in `headers` — `x-api-key` @@ -856,9 +879,10 @@ impl SamplingClient { sent_bearer, ); let server_message = user_facing_api_error_message(status, bytes.as_ref()); - return Err(SamplingError::Auth(format!( - "Unauthorized (401): {server_message}" - ))); + return Err(auth_rejected( + format!("Unauthorized (401): {server_message}"), + sent_bearer, + )); } let message = user_facing_api_error_message(status, bytes.as_ref()); return Err(SamplingError::Api { @@ -911,7 +935,10 @@ impl SamplingClient { deployment_id: payload.x_grok_deployment_id.as_deref(), user_id: payload.x_grok_user_id.as_deref(), }; - let (builder, sent_bearer) = self.post(self.endpoint("chat/completions")); + let SentRequest { + builder, + sent_bearer, + } = self.post(self.endpoint("chat/completions")); let http_request = grok_headers.apply(builder).json(&payload); let response = http_request.send().await.map_err(|e| { @@ -968,7 +995,10 @@ impl SamplingClient { deployment_id: payload.x_grok_deployment_id.as_deref(), user_id: payload.x_grok_user_id.as_deref(), }; - let (builder, sent_bearer) = self.post(self.endpoint("chat/completions")); + let SentRequest { + builder, + sent_bearer, + } = self.post(self.endpoint("chat/completions")); let http_request = grok_headers .apply(builder) .header(ACCEPT, HeaderValue::from_static("text/event-stream")) @@ -1009,9 +1039,10 @@ impl SamplingClient { let endpoint = self.endpoint("chat/completions"); let body = response.bytes().await.unwrap_or_default(); let server_message = user_facing_api_error_message(status, body.as_ref()); - return Err(SamplingError::Auth(format!( - "Unauthorized (401) from {endpoint}: {server_message}" - ))); + return Err(auth_rejected( + format!("Unauthorized (401) from {endpoint}: {server_message}"), + sent_bearer.as_deref(), + )); } let bytes = response.bytes().await?; @@ -1183,7 +1214,10 @@ impl SamplingClient { // it in post-serialize. This is the last surviving piece of the // old raw_output machinery. xai_grok_sampling_types::patch_reasoning_text_types(&mut request_body); - let (builder, sent_bearer) = self.post(self.endpoint("responses")); + let SentRequest { + builder, + sent_bearer, + } = self.post(self.endpoint("responses")); let http_request = grok_headers.apply(builder).json(&request_body); let response = http_request.send().await.map_err(|e| { @@ -1205,9 +1239,10 @@ impl SamplingClient { ); let endpoint = self.endpoint("responses"); let server_message = user_facing_api_error_message(status, bytes.as_ref()); - return Err(SamplingError::Auth(format!( - "Unauthorized (401) from {endpoint}: {server_message}" - ))); + return Err(auth_rejected( + format!("Unauthorized (401) from {endpoint}: {server_message}"), + sent_bearer.as_deref(), + )); } let message = user_facing_api_error_message(status, bytes.as_ref()); @@ -1327,7 +1362,10 @@ impl SamplingClient { .defaults .doom_loop_recovery .map(crate::doom_loop::DoomLoopSignalCollector::new); - let (builder, sent_bearer) = self.post(self.endpoint("responses")); + let SentRequest { + builder, + sent_bearer, + } = self.post(self.endpoint("responses")); let mut http_request = grok_headers .apply(builder) .header(ACCEPT, HeaderValue::from_static("text/event-stream")); @@ -1369,9 +1407,10 @@ impl SamplingClient { let endpoint = self.endpoint("responses"); let body = response.bytes().await.unwrap_or_default(); let server_message = user_facing_api_error_message(status, body.as_ref()); - return Err(SamplingError::Auth(format!( - "Unauthorized (401) from {endpoint}: {server_message}" - ))); + return Err(auth_rejected( + format!("Unauthorized (401) from {endpoint}: {server_message}"), + sent_bearer.as_deref(), + )); } let model_metadata = extract_model_metadata(response.headers()); let retry_after_secs = extract_retry_after(response.headers()); @@ -1528,7 +1567,10 @@ impl SamplingClient { deployment_id: request.x_grok_deployment_id.as_deref(), user_id: request.x_grok_user_id.as_deref(), }; - let (builder, sent_bearer) = self.post(self.endpoint("messages")); + let SentRequest { + builder, + sent_bearer, + } = self.post(self.endpoint("messages")); let http_request = grok_headers.apply(builder).json(&request.inner); let response = http_request.send().await.map_err(|e| { @@ -1550,9 +1592,10 @@ impl SamplingClient { ); let endpoint = self.endpoint("messages"); let server_message = user_facing_api_error_message(status, bytes.as_ref()); - return Err(SamplingError::Auth(format!( - "Unauthorized (401) from {endpoint}: {server_message}" - ))); + return Err(auth_rejected( + format!("Unauthorized (401) from {endpoint}: {server_message}"), + sent_bearer.as_deref(), + )); } let message = user_facing_api_error_message(status, bytes.as_ref()); @@ -1637,7 +1680,10 @@ impl SamplingClient { deployment_id: request.x_grok_deployment_id.as_deref(), user_id: request.x_grok_user_id.as_deref(), }; - let (builder, sent_bearer) = self.post(self.endpoint("messages")); + let SentRequest { + builder, + sent_bearer, + } = self.post(self.endpoint("messages")); let http_request = grok_headers .apply(builder) .header(ACCEPT, HeaderValue::from_static("text/event-stream")) @@ -1675,9 +1721,10 @@ impl SamplingClient { let endpoint = self.endpoint("messages"); let body = response.bytes().await.unwrap_or_default(); let server_message = user_facing_api_error_message(status, body.as_ref()); - return Err(SamplingError::Auth(format!( - "Unauthorized (401) from {endpoint}: {server_message}" - ))); + return Err(auth_rejected( + format!("Unauthorized (401) from {endpoint}: {server_message}"), + sent_bearer.as_deref(), + )); } let model_metadata = extract_model_metadata(response.headers()); let retry_after_secs = extract_retry_after(response.headers()); @@ -2323,7 +2370,7 @@ mod tests { let mut config = minimal_config(); config.header_injector = Some(std::sync::Arc::new(TestInjector)); let client = SamplingClient::new(config).expect("build"); - let (builder, _sent) = client.post("http://localhost/test"); + let SentRequest { builder, .. } = client.post("http://localhost/test"); let req = builder.build().expect("build request"); assert!( req.headers().contains_key("traceparent"), @@ -2403,7 +2450,10 @@ mod tests { ..minimal_config() }; let client = SamplingClient::new(cfg).expect("client should build"); - let (_builder, bearer) = client.post("https://example.test/v1/chat/completions"); + let SentRequest { + sent_bearer: bearer, + .. + } = client.post("https://example.test/v1/chat/completions"); assert_eq!(bearer.as_deref(), Some("r-1234567890")); assert_eq!( bearer.as_deref().map(str::len), @@ -2422,7 +2472,10 @@ mod tests { ..minimal_config() }; let client = SamplingClient::new(cfg).expect("client should build"); - let (_builder, bearer) = client.post("https://example.test/v1/messages"); + let SentRequest { + sent_bearer: bearer, + .. + } = client.post("https://example.test/v1/messages"); assert_eq!(bearer.as_deref(), Some("c-key-abc123")); assert_eq!( bearer.as_deref().map(str::len), @@ -2439,7 +2492,10 @@ mod tests { ..minimal_config() }; let client = SamplingClient::new(cfg).expect("client should build"); - let (_builder, bearer) = client.post("https://example.test/v1/chat/completions"); + let SentRequest { + sent_bearer: bearer, + .. + } = client.post("https://example.test/v1/chat/completions"); assert!(bearer.is_none()); } @@ -2468,7 +2524,10 @@ mod tests { }; let client = SamplingClient::new(cfg).expect("client should build"); - let (_builder, sent_at_build) = client.post("https://example.test/v1/responses"); + let SentRequest { + sent_bearer: sent_at_build, + .. + } = client.post("https://example.test/v1/responses"); // The 401 kicks recovery; the resolver rotates before the callback runs. *resolver.0.lock().unwrap() = "fresh-token-newtail99".to_string(); @@ -2496,7 +2555,7 @@ mod tests { ..minimal_config() }; let client = SamplingClient::new(cfg).expect("client should build"); - let (builder, _sent) = client.post("https://example.test/v1/messages"); + let SentRequest { builder, .. } = client.post("https://example.test/v1/messages"); let request = builder.build().expect("request should build"); let auth = request .headers() @@ -2522,7 +2581,7 @@ mod tests { ..minimal_config() }; let client = SamplingClient::new(cfg).expect("client should build"); - let (builder, _sent) = client.post("https://example.test/v1/responses"); + let SentRequest { builder, .. } = client.post("https://example.test/v1/responses"); let request = builder.build().expect("request should build"); let auth_count = request.headers().get_all(AUTHORIZATION).iter().count(); assert_eq!( @@ -2548,7 +2607,7 @@ mod tests { ..minimal_config() }; let client = SamplingClient::new(cfg).expect("client should build"); - let (builder, _sent) = client.post("https://example.test/v1/messages"); + let SentRequest { builder, .. } = client.post("https://example.test/v1/messages"); let request = builder.build().expect("request should build"); let api_key = request .headers() @@ -2572,7 +2631,8 @@ mod tests { ..minimal_config() }; let client = SamplingClient::new(cfg).expect("client should build"); - let (_builder, sent_bearer) = client.post("https://example.test/v1/chat/completions"); + let SentRequest { sent_bearer, .. } = + client.post("https://example.test/v1/chat/completions"); client.record_401_attribution( crate::attribution::SamplingConsumer::ChatCompletionsStream, sent_bearer.as_deref(), @@ -2636,7 +2696,10 @@ mod tests { ..minimal_config() }; let client = SamplingClient::new(cfg).expect("client should build"); - let (builder, sent) = client.post("https://example.test/v1/responses"); + let SentRequest { + builder, + sent_bearer: sent, + } = client.post("https://example.test/v1/responses"); let request = builder.body("").build().expect("request should build"); assert_eq!(sent, None, "capture must agree: nothing was sent"); assert!( @@ -2670,7 +2733,7 @@ mod tests { let client = SamplingClient::new(cfg).expect("client should build"); // Build a request to inspect the final headers. - let (builder, _sent) = client.post("https://example.test/v1/responses"); + let SentRequest { builder, .. } = client.post("https://example.test/v1/responses"); let request = builder.body("").build().expect("request should build"); let auth_values: Vec<_> = request.headers().get_all(AUTHORIZATION).iter().collect(); diff --git a/crates/codegen/xai-grok-sampler/src/events.rs b/crates/codegen/xai-grok-sampler/src/events.rs index a242de9..b65e841 100644 --- a/crates/codegen/xai-grok-sampler/src/events.rs +++ b/crates/codegen/xai-grok-sampler/src/events.rs @@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize}; use xai_grok_sampling_types::{ ConversationResponse, EmptyResponseContext, ResponseModelMetadata, SamplingError, + SentCredential, }; use crate::metrics::InferenceLatencyStats; @@ -168,6 +169,11 @@ pub struct SamplingErrorInfo { /// Telemetry only; `None` for terminal-response detections. #[serde(default, skip_serializing_if = "Option::is_none")] pub doom_loop_aborted_at_chunk: Option, + /// Meaningful only when `kind == Auth`: whether the rejected request + /// actually carried a credential on the wire. Defaults to `Unknown` + /// (charge-the-budget behavior) for payloads from older peers. + #[serde(default, skip_serializing_if = "SentCredential::is_unknown")] + pub credential: SentCredential, } /// Coarse-grained classification of a sampling failure. @@ -217,7 +223,7 @@ impl From<&SamplingError> for SamplingErrorInfo { let message = err.to_string(); let (kind, status_code, retry_after_secs, model_metadata) = match err { - SamplingError::Auth(_) => (SamplingErrorKind::Auth, None, None, None), + SamplingError::Auth { .. } => (SamplingErrorKind::Auth, None, None, None), SamplingError::InvalidConfiguration(_) => (SamplingErrorKind::Api, None, None, None), SamplingError::Http(_) => (SamplingErrorKind::Http, None, None, None), SamplingError::Serialization(_) => (SamplingErrorKind::Serialization, None, None, None), @@ -264,6 +270,10 @@ impl From<&SamplingError> for SamplingErrorInfo { } => (Some(triggers.clone()), *aborted_at_chunk), _ => (None, None), }; + let credential = match err { + SamplingError::Auth { credential, .. } => *credential, + _ => SentCredential::Unknown, + }; Self { kind, @@ -275,6 +285,7 @@ impl From<&SamplingError> for SamplingErrorInfo { empty_response_context, doom_loop_triggers, doom_loop_aborted_at_chunk, + credential, } } } @@ -286,7 +297,7 @@ mod tests { #[test] fn auth_variant_classified_as_auth() { - let err = SamplingError::Auth("bad token".into()); + let err = SamplingError::auth_unknown("bad token"); let info = SamplingErrorInfo::from(&err); assert_eq!(info.kind, SamplingErrorKind::Auth); assert_eq!(info.status_code, None); @@ -296,6 +307,18 @@ mod tests { assert!(info.message.contains("bad token")); } + /// A payload from a peer that predates `credential` must still parse, + /// defaulting to `Unknown` (charge-the-budget behavior). + #[test] + fn info_without_credential_field_deserializes_to_unknown() { + let info: SamplingErrorInfo = serde_json::from_str( + r#"{"kind":"Auth","status_code":401,"message":"x","is_retryable":false, + "retry_after_secs":null,"model_metadata":null}"#, + ) + .unwrap(); + assert_eq!(info.credential, SentCredential::Unknown); + } + #[test] fn invalid_configuration_classified_as_api() { let err = SamplingError::InvalidConfiguration("missing model"); diff --git a/crates/codegen/xai-grok-sampler/src/handle.rs b/crates/codegen/xai-grok-sampler/src/handle.rs index 1b2fe7f..f9c69b0 100644 --- a/crates/codegen/xai-grok-sampler/src/handle.rs +++ b/crates/codegen/xai-grok-sampler/src/handle.rs @@ -149,8 +149,8 @@ impl SamplerHandle { request_id: cancel_id, }); completion_rx.await.unwrap_or_else(|_| { - Err(SamplingError::Auth( - "sampler actor dropped before completion".to_string(), + Err(SamplingError::auth_unknown( + "sampler actor dropped before completion", )) }) } diff --git a/crates/codegen/xai-grok-sampler/src/retry.rs b/crates/codegen/xai-grok-sampler/src/retry.rs index 5f75239..071bbd3 100644 --- a/crates/codegen/xai-grok-sampler/src/retry.rs +++ b/crates/codegen/xai-grok-sampler/src/retry.rs @@ -173,26 +173,20 @@ 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. + // Shared retry vetoes (`SamplingError::is_retry_vetoed`, also used by + // one-shot callers like /btw): + // - x-should-retry: false — trust the server, it knows if the error is + // request-content-caused (e.g. malformed tool call in history) vs + // transient. x-should-retry: true is intentionally NOT handled — the + // header only suppresses retries; forcing them on non-retryable + // statuses could amplify failures. + // - Context-window / size overflow — deterministic, re-sending the same + // (or larger) payload always fails, whatever status the backend used. // // 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). - if err.is_context_length_error() { + if err.is_retry_vetoed() { return RetryDecision::Fatal(clone_error(err)); } @@ -263,10 +257,10 @@ pub fn format_sampling_error(err: &SamplingError, retry_count: Option) -> S }; match err { - SamplingError::Auth(msg) => { + SamplingError::Auth { message, .. } => { format!( "{}Authentication failed: {}. Please check your API key configuration.", - retry_prefix, msg + retry_prefix, message ) } SamplingError::InvalidConfiguration(msg) => { @@ -389,7 +383,13 @@ pub fn format_sampling_error(err: &SamplingError, retry_count: Option) -> S /// retry budget re-generating a response that fails the same way. pub(crate) fn clone_error(err: &SamplingError) -> SamplingError { match err { - SamplingError::Auth(msg) => SamplingError::Auth(msg.clone()), + SamplingError::Auth { + message, + credential, + } => SamplingError::Auth { + message: message.clone(), + credential: *credential, + }, SamplingError::InvalidConfiguration(msg) => SamplingError::InvalidConfiguration(msg), SamplingError::Http(e) => { // reqwest::Error is not Clone; preserve the rendered message @@ -520,9 +520,9 @@ mod tests { #[test] fn classify_auth_error_emits_to_session() { - let err = SamplingError::Auth("bad token".into()); + let err = SamplingError::auth_unknown("bad token"); match classify_error(&err, 0, 5, RATE_LIMIT_RETRY_THRESHOLD) { - RetryDecision::EmitToSession(SamplingError::Auth(_)) => {} + RetryDecision::EmitToSession(SamplingError::Auth { .. }) => {} other => panic!("expected EmitToSession(Auth), got {other:?}"), } } @@ -763,14 +763,14 @@ mod tests { #[test] fn format_includes_retry_prefix_when_count_present() { - let err = SamplingError::Auth("bad".into()); + let err = SamplingError::auth_unknown("bad"); let s = format_sampling_error(&err, Some(3)); assert!(s.starts_with("Request failed after 3 retries.")); } #[test] fn format_omits_retry_prefix_when_count_absent() { - let err = SamplingError::Auth("bad".into()); + let err = SamplingError::auth_unknown("bad"); let s = format_sampling_error(&err, None); assert!(!s.starts_with("Request failed after")); assert!(s.starts_with("Authentication failed:")); diff --git a/crates/codegen/xai-grok-sampler/src/stream/collect.rs b/crates/codegen/xai-grok-sampler/src/stream/collect.rs index 96d83bd..214f4db 100644 --- a/crates/codegen/xai-grok-sampler/src/stream/collect.rs +++ b/crates/codegen/xai-grok-sampler/src/stream/collect.rs @@ -53,6 +53,7 @@ pub async fn collect_response( empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: xai_grok_sampling_types::SentCredential::Unknown, }) } diff --git a/crates/codegen/xai-grok-sampling-types/src/error.rs b/crates/codegen/xai-grok-sampling-types/src/error.rs index 1407754..f18db87 100644 --- a/crates/codegen/xai-grok-sampling-types/src/error.rs +++ b/crates/codegen/xai-grok-sampling-types/src/error.rs @@ -76,6 +76,65 @@ pub struct ResponseModelMetadata { pub models_etag: Option, } +/// Wire-credential provenance of a request that failed authentication. +/// +/// A 401 for a request that went out with **no** credential header (a +/// fail-closed send while the bearer resolver had nothing wire-valid) is +/// not evidence against the credential itself; retry policies use this to +/// avoid charging credential-rejection budgets for such sends. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +#[non_exhaustive] +pub enum SentCredential { + /// The request carried a credential; the server rejected it. + Sent, + /// The request went out with no credential header at all. + Missing, + /// Provenance unknown (synthesized or legacy errors). Retry policies + /// treat this like [`SentCredential::Sent`] — fail closed toward + /// terminating rather than retrying forever. + #[default] + Unknown, +} + +/// Hand-written so an unrecognized value from a newer peer degrades to +/// `Unknown` instead of failing the whole containing payload +/// (`#[serde(other)]` is not available on externally-tagged enums). +impl<'de> Deserialize<'de> for SentCredential { + fn deserialize>( + deserializer: D, + ) -> std::result::Result { + Ok( + match std::borrow::Cow::::deserialize(deserializer)?.as_ref() { + "sent" => Self::Sent, + "missing" => Self::Missing, + _ => Self::Unknown, + }, + ) + } +} + +impl SentCredential { + /// Classify from the credential fragment captured when the request was + /// built (`None` = no credential header was stamped on the wire). + pub fn from_sent_fragment(fragment: Option<&str>) -> Self { + if fragment.is_some() { + Self::Sent + } else { + Self::Missing + } + } + + pub fn is_missing(self) -> bool { + matches!(self, Self::Missing) + } + + /// By reference so it can serve as a serde `skip_serializing_if`. + pub fn is_unknown(&self) -> bool { + matches!(self, Self::Unknown) + } +} + /// Display prefix of [`SamplingError::Serialization`]. Shared with the /// variant's `#[error(...)]` template so [`SamplingError::serialization_from_rendered`] /// can never drift from what Display actually emits. @@ -83,8 +142,12 @@ const SERIALIZATION_DISPLAY_PREFIX: &str = "serialization error: "; #[derive(Debug, Error)] pub enum SamplingError { - #[error("{0}")] - Auth(String), + #[error("{message}")] + Auth { + message: String, + /// Whether the rejected request actually carried a credential. + credential: SentCredential, + }, #[error("invalid client configuration: {0}")] InvalidConfiguration(&'static str), #[error("request error: {0}")] @@ -132,6 +195,16 @@ pub enum SamplingError { } impl SamplingError { + /// Auth error of unknown wire provenance — for paths that never sent a + /// request (config validation, cancellation, actor teardown) or that + /// lost the provenance (legacy round trips). + pub fn auth_unknown(message: impl Into) -> Self { + Self::Auth { + message: message.into(), + credential: SentCredential::Unknown, + } + } + /// Rebuild a `Serialization` error from a rendered message for non-`Clone` /// contexts; it must stay `Serialization` so it remains non-retryable. pub fn serialization_message(msg: impl fmt::Display) -> Self { @@ -161,7 +234,7 @@ impl SamplingError { // can race with invalid_grant_threshold to wipe auth.json. matches!( self, - SamplingError::Auth(_) + SamplingError::Auth { .. } | SamplingError::Api { status: StatusCode::UNAUTHORIZED, .. @@ -239,12 +312,12 @@ impl SamplingError { pub fn is_retryable(&self) -> bool { match self { - SamplingError::Auth(_) => false, + SamplingError::Auth { .. } => false, SamplingError::InvalidConfiguration(_) => false, SamplingError::Http(err) => is_retryable_reqwest(err), SamplingError::Serialization(_) => false, SamplingError::Api { status, .. } => { - matches!(status.as_u16(), 429 | 500 | 502 | 503 | 504 | 520) + matches!(status.as_u16(), 429 | 500 | 502 | 503 | 504 | 520 | 529) } SamplingError::EventStreamError(_) => true, SamplingError::StreamError { .. } => true, @@ -289,6 +362,42 @@ impl SamplingError { _ => false, } } + + /// Capacity / overload: HTTP 529, a 5xx whose message clearly says + /// overloaded (proxies wrap stream overloads in a 500), or a stream + /// error whose parsed `error_type` is a capacity type (`overloaded_error` + /// / `service_unavailable_error`). Never reachable from a 4xx or a + /// request-shaped stream error, whatever the message text. Transient — + /// worth a short, bounded retry at the call site. + pub fn is_overloaded(&self) -> bool { + match self { + SamplingError::Api { + status, message, .. + } => { + status.as_u16() == 529 + || (status.is_server_error() && message_looks_overloaded(message)) + } + // `error_type` is already parsed from the stream payload — trust + // it alone; matching message text here would let a request-shaped + // error that merely mentions "overloaded" retry. + SamplingError::StreamError { error_type, .. } => { + error_type.eq_ignore_ascii_case("overloaded_error") + || error_type.eq_ignore_ascii_case("service_unavailable_error") + } + _ => false, + } + } + + /// Retry vetoes shared by every retry loop — the sampler actor's + /// `classify_error` and one-shot callers like `/btw`. One definition so + /// a new veto lands everywhere at once: + /// - `x-should-retry: false` — the server says the failure is + /// request-content-caused, not transient. + /// - Context-length overflow — deterministic; re-sending the same + /// payload always fails. + pub fn is_retry_vetoed(&self) -> bool { + self.should_retry_header() == Some(false) || self.is_context_length_error() + } } impl From for SamplingError { @@ -430,6 +539,7 @@ pub fn is_context_length_error(message: &str) -> bool { || m.contains("maximum prompt length") || m.contains("maximum context length") || m.contains("context_length_exceeded") + || (m.contains("current message") && m.contains("exceeds budget")) } /// Decide whether a [`reqwest::Error`] is worth retrying. @@ -452,10 +562,164 @@ pub fn is_retryable_reqwest(err: &reqwest::Error) -> bool { false } +/// Capacity-style provider text: "Overloaded" / `overloaded_error` (possibly +/// proxy-wrapped) or `service_unavailable_error` (503-shaped capacity). +fn message_looks_overloaded(message: &str) -> bool { + let m = message.to_ascii_lowercase(); + m.contains("overloaded") || m.contains("service_unavailable_error") +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn overloaded_detects_stream_and_api_shapes() { + assert!( + SamplingError::StreamError { + error_type: "overloaded_error".into(), + message: "Overloaded".into(), + } + .is_overloaded() + ); + assert!( + SamplingError::Api { + status: StatusCode::INTERNAL_SERVER_ERROR, + message: "stream error (overloaded_error): Overloaded".into(), + model_metadata: None, + retry_after_secs: None, + should_retry: None, + } + .is_overloaded() + ); + assert!( + SamplingError::Api { + status: StatusCode::from_u16(529).unwrap(), + message: "capacity".into(), + model_metadata: None, + retry_after_secs: None, + should_retry: None, + } + .is_overloaded() + ); + assert!( + SamplingError::Api { + status: StatusCode::from_u16(529).unwrap(), + message: "capacity".into(), + model_metadata: None, + retry_after_secs: None, + should_retry: None, + } + .is_retryable() + ); + assert!(!SamplingError::auth_unknown("nope").is_overloaded()); + assert!( + !SamplingError::Api { + status: StatusCode::BAD_REQUEST, + message: "invalid json".into(), + model_metadata: None, + retry_after_secs: None, + should_retry: None, + } + .is_overloaded() + ); + // Only server errors classify on message text — a 4xx that merely + // mentions "overloaded" is a request error, not capacity. + assert!( + !SamplingError::Api { + status: StatusCode::BAD_REQUEST, + message: "field `overloaded` is not a valid parameter".into(), + model_metadata: None, + retry_after_secs: None, + should_retry: None, + } + .is_overloaded() + ); + // Stream errors classify on the parsed error_type only — a + // request-shaped stream error mentioning "overloaded" is not capacity. + assert!( + !SamplingError::StreamError { + error_type: "invalid_request_error".into(), + message: "tool result mentions overloaded".into(), + } + .is_overloaded() + ); + assert!( + SamplingError::StreamError { + error_type: "service_unavailable_error".into(), + message: "upstream capacity".into(), + } + .is_overloaded() + ); + } + + #[test] + fn overloaded_message_matches_backend_variants() { + // 5xx messages that classify as capacity. + for msg in [ + "Overloaded", + "stream error (overloaded_error): Overloaded", + "overloaded_error", + "service_unavailable_error: try again", + ] { + assert!( + SamplingError::Api { + status: StatusCode::INTERNAL_SERVER_ERROR, + message: msg.into(), + model_metadata: None, + retry_after_secs: None, + should_retry: None, + } + .is_overloaded(), + "expected overloaded for message: {msg}" + ); + } + // 5xx messages that do not. + for msg in ["upstream connect timeout", "internal error"] { + assert!( + !SamplingError::Api { + status: StatusCode::INTERNAL_SERVER_ERROR, + message: msg.into(), + model_metadata: None, + retry_after_secs: None, + should_retry: None, + } + .is_overloaded(), + "expected not overloaded for message: {msg}" + ); + } + } + + #[test] + fn retry_veto_covers_header_and_context_length() { + let vetoed_by_header = SamplingError::Api { + status: StatusCode::from_u16(529).unwrap(), + message: "capacity".into(), + model_metadata: None, + retry_after_secs: None, + should_retry: Some(false), + }; + assert!(vetoed_by_header.is_retry_vetoed()); + + let vetoed_by_context = SamplingError::Api { + status: StatusCode::from_u16(529).unwrap(), + message: "prompt is too long: 300000 tokens > 200000 maximum".into(), + model_metadata: None, + retry_after_secs: None, + should_retry: None, + }; + assert!(vetoed_by_context.is_retry_vetoed()); + + let not_vetoed = SamplingError::Api { + status: StatusCode::from_u16(529).unwrap(), + message: "capacity".into(), + model_metadata: None, + retry_after_secs: None, + should_retry: None, + }; + assert!(!not_vetoed.is_retry_vetoed()); + } + #[test] fn context_length_error_matches_backend_variants() { for msg in [ @@ -465,10 +729,20 @@ mod tests { "This model's maximum context length is 200000 tokens", "invalid_request_error: prompt is too long: 300000 tokens > 200000 maximum", "error type: context_length_exceeded", + "Failed to start sampling: [conversation] Current message (1000000 tokens) exceeds budget (500000 tokens)", + "API error (status 400 Bad Request): invalid-argument: Failed to start sampling: [conversation] Current message (1000000 tokens) exceeds budget (500000 tokens)", + "compact failed: API error (status 400 Bad Request): invalid-argument: Failed to start sampling: [conversation] Current message (1000000 tokens) exceeds budget (500000 tokens)", + "Current message (600000) exceeds budget (500000)", ] { assert!(is_context_length_error(msg), "should match: {msg}"); } - for msg in ["rate limited", "internal server error", "connection reset"] { + for msg in [ + "rate limited", + "internal server error", + "connection reset", + "Attached file content (300000 tokens) causes message to exceed budget", + "compact index estimate 2.0 GB exceeds budget 1.0 GB", + ] { assert!(!is_context_length_error(msg), "should not match: {msg}"); } // The method delegates for the Api/StreamError variants. @@ -487,7 +761,7 @@ mod tests { } .is_context_length_error() ); - assert!(!SamplingError::Auth("nope".into()).is_context_length_error()); + assert!(!SamplingError::auth_unknown("nope").is_context_length_error()); } #[test] @@ -671,10 +945,31 @@ mod tests { #[test] fn auth_variant_is_auth_error() { - let err = SamplingError::Auth("bad key".into()); + let err = SamplingError::auth_unknown("bad key"); assert!(err.is_auth_error()); } + /// Known values round-trip; an unrecognized value from a newer peer + /// degrades to `Unknown` instead of failing the containing payload. + #[test] + fn sent_credential_wire_compat() { + for (json, expected) in [ + ("\"sent\"", SentCredential::Sent), + ("\"missing\"", SentCredential::Missing), + ("\"unknown\"", SentCredential::Unknown), + ("\"some-future-variant\"", SentCredential::Unknown), + ] { + assert_eq!( + serde_json::from_str::(json).unwrap(), + expected + ); + } + assert_eq!( + serde_json::to_string(&SentCredential::Missing).unwrap(), + "\"missing\"" + ); + } + #[test] fn rate_limited_api_error_is_detected() { let err = SamplingError::Api { @@ -701,7 +996,7 @@ mod tests { }; assert!(!server_error.is_rate_limited()); - let auth_error = SamplingError::Auth("bad key".into()); + let auth_error = SamplingError::auth_unknown("bad key"); assert!(!auth_error.is_rate_limited()); let timeout = SamplingError::IdleTimeout { elapsed_secs: 30 }; @@ -734,7 +1029,7 @@ mod tests { #[test] fn retry_after_returns_none_for_non_api_errors() { - assert_eq!(SamplingError::Auth("x".into()).retry_after(), None); + assert_eq!(SamplingError::auth_unknown("x").retry_after(), None); assert_eq!( SamplingError::IdleTimeout { elapsed_secs: 10 }.retry_after(), None diff --git a/crates/codegen/xai-grok-sampling-types/src/lib.rs b/crates/codegen/xai-grok-sampling-types/src/lib.rs index b2fca86..5414b68 100644 --- a/crates/codegen/xai-grok-sampling-types/src/lib.rs +++ b/crates/codegen/xai-grok-sampling-types/src/lib.rs @@ -21,7 +21,7 @@ pub use self::doom_loop::{ }; pub use self::error::{ EmptyReason, EmptyResponseContext, ResponseModelMetadata, Result, SamplingError, - is_context_length_error, status_user_message, user_facing_api_error_message, + SentCredential, is_context_length_error, status_user_message, user_facing_api_error_message, }; pub use self::tool_overrides::{ ClearableField, SearchDateBound, SearchDateBoundError, ToolOverrides, ToolOverridesUpdate, diff --git a/crates/codegen/xai-grok-shell-base/Cargo.toml b/crates/codegen/xai-grok-shell-base/Cargo.toml index b2340bf..6a81e69 100644 --- a/crates/codegen/xai-grok-shell-base/Cargo.toml +++ b/crates/codegen/xai-grok-shell-base/Cargo.toml @@ -6,6 +6,8 @@ edition.workspace = true description = "Foundation modules for the grok shell crate family: environment presets, CPU profiling, and process/filesystem utilities." [features] +# Forwarded from shell; empty here (no local-ws code in this crate). +local-workspace = [] # Exposes `#[cfg(test)]`-only helpers (`env::EnvVarGuard`, # `cpu_profile` test seams) to downstream crates' test targets. Enabled by test-support = ["xai-grok-env/test-support"] diff --git a/crates/codegen/xai-grok-shell/CHANGELOG.md b/crates/codegen/xai-grok-shell/CHANGELOG.md index a6aa715..3873099 100644 --- a/crates/codegen/xai-grok-shell/CHANGELOG.md +++ b/crates/codegen/xai-grok-shell/CHANGELOG.md @@ -1,5 +1,24 @@ # Changelog +# 0.2.117 — 2026-07-30 + +## Features + +- **GROK_EXTRA_CA_BUNDLE** env var allows adding custom TLS root certificates. + +## Bug Fixes + +- **Stop command** now terminates all background subagents from prior turns. +- **kill_task** tool now correctly reports when a task does not exist over ACP connections. +- **get_task_output** no longer waits the full timeout for already-finished tasks over ACP. +- **/usage** command and billing UI are hidden for enterprise auth setups. +- **Plan approval** no longer starts Build when pressing Enter without notes in revise mode. + +## Performance + +- **Terminal resize** is much faster on long conversations in fullscreen mode. + + # 0.2.116 — 2026-07-30 ## Features diff --git a/crates/codegen/xai-grok-shell/Cargo.toml b/crates/codegen/xai-grok-shell/Cargo.toml index e793778..b54ca55 100644 --- a/crates/codegen/xai-grok-shell/Cargo.toml +++ b/crates/codegen/xai-grok-shell/Cargo.toml @@ -1,7 +1,7 @@ [package] license = "Apache-2.0" name = "xai-grok-shell" -version = "0.2.116" +version = "0.2.117" edition.workspace = true [features] @@ -12,7 +12,12 @@ dhat-heap = ["dep:dhat"] # load, and bench tests. Off by default; the tests/benches that use it declare # it via `required-features`. test-support = [] -default-bazel = ["test-support"] +# Local Computer Hub workspace_server (own/attach + crash-restart). Requires +local-workspace = [] +default-bazel = [ + "local-workspace", + "test-support", +] [dependencies] dunce = { workspace = true } @@ -189,6 +194,8 @@ windows = { workspace = true } [dev-dependencies] criterion = { workspace = true } +# Pre-main unified-log redirect (`test_support::redirect_unified_log_for_tests`). +ctor = { workspace = true } filetime = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } @@ -223,6 +230,10 @@ name = "fork_copy" harness = false required-features = ["test-support"] +[[bench]] +name = "skills_watcher_startup" +harness = false + [[test]] name = "test_leader_soak" required-features = ["test-support"] diff --git a/crates/codegen/xai-grok-shell/benches/skills_watcher_startup.rs b/crates/codegen/xai-grok-shell/benches/skills_watcher_startup.rs new file mode 100644 index 0000000..34efac3 --- /dev/null +++ b/crates/codegen/xai-grok-shell/benches/skills_watcher_startup.rs @@ -0,0 +1,150 @@ +//! Skills file-watcher startup latency. +//! +//! Times OS watch registration for a project-tier `.claude` tree with a large +//! `worktrees/` subtree (Bazel-like fan-out). Compares: +//! +//! - **scoped** — current `SkillsFileWatcher::start_with_dirs` (vendor root +//! non-recursive + skills/commands/workflows only) +//! - **recursive_control** — full `RecursiveMode::Recursive` on `.claude` +//! (pre-fix project-tier behavior on Linux: one inotify wd per directory) +//! +//! Fixture sizes stay comparable across scenarios. Medians land under +//! `target/criterion/skills_watcher_startup/`. +//! +//! ```text +//! cargo bench -p xai-grok-shell --bench skills_watcher_startup +//! # optional scale: +//! GROK_SKILLS_WATCHER_BENCH_DIRS=12000 cargo bench -p xai-grok-shell --bench skills_watcher_startup +//! ``` +//! +//! On macOS, recursive FSEvents is cheap so both arms may be close. On Linux +//! inotify, `recursive_control` scales with directory count; `scoped` stays flat. + +use std::fs; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use criterion::{BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use notify::RecursiveMode; +use notify_debouncer_mini::new_debouncer; +use tempfile::TempDir; +use xai_grok_shell::config::watcher::SkillsFileWatcher; + +/// Default dirs under `.claude/worktrees/` (override with env). +const DEFAULT_WORKTREE_DIRS: usize = 6_000; + +fn worktree_dir_count() -> usize { + std::env::var("GROK_SKILLS_WATCHER_BENCH_DIRS") + .ok() + .and_then(|v| v.parse().ok()) + .filter(|&n| n > 0) + .unwrap_or(DEFAULT_WORKTREE_DIRS) +} + +/// Nested groups of 100 so the tree has width and depth. +fn make_nested_dirs(base: &Path, count: usize) { + for i in 0..count { + let dir = base.join(format!("g{}", i / 100)).join(format!("d{i}")); + fs::create_dir_all(&dir).unwrap(); + } +} + +struct Fixture { + _root: TempDir, + project: PathBuf, + claude: PathBuf, + grok_home: PathBuf, +} + +/// Project with a real skill and a fat `.claude/worktrees` tree. +fn build_fixture(worktree_dirs: usize) -> Fixture { + let root = TempDir::new().unwrap(); + let project = root.path().join("project"); + let claude = project.join(".claude"); + let skills = claude.join("skills").join("alpha"); + fs::create_dir_all(&skills).unwrap(); + fs::write(skills.join("SKILL.md"), "# alpha\n").unwrap(); + + let worktrees = claude.join("worktrees").join("wt1"); + make_nested_dirs(&worktrees, worktree_dirs); + + let grok_home = root.path().join("grok-home"); + fs::create_dir_all(&grok_home).unwrap(); + + Fixture { + _root: root, + project, + claude, + grok_home, + } +} + +fn start_scoped(fixture: &Fixture) -> SkillsFileWatcher { + let dirs = vec![fixture.claude.clone()]; + let (watcher, _rx) = SkillsFileWatcher::start_with_dirs( + &dirs, + &fixture.grok_home, + Some(fixture.project.as_path()), + ) + .expect("scoped skills watcher should start"); + watcher +} + +/// Pre-fix control: one recursive watch on the whole project `.claude`. +fn start_recursive_control( + claude: &Path, +) -> notify_debouncer_mini::Debouncer { + let mut debouncer = new_debouncer(Duration::from_secs(2), |_| {}).expect("debouncer"); + debouncer + .watcher() + .watch(claude, RecursiveMode::Recursive) + .expect("recursive watch"); + debouncer +} + +fn bench_skills_watcher_startup(c: &mut Criterion) { + let n = worktree_dir_count(); + let fixture = build_fixture(n); + + eprintln!( + "skills_watcher_startup fixture: project={:?} worktree_dirs={n}", + fixture.project + ); + + let mut group = c.benchmark_group("skills_watcher_startup"); + group.sample_size(20); + group.throughput(Throughput::Elements(n as u64)); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(8)); + + group.bench_function(BenchmarkId::new("scoped", n), |b| { + b.iter_batched(|| (), |()| start_scoped(&fixture), BatchSize::PerIteration); + }); + + group.bench_function(BenchmarkId::new("recursive_control", n), |b| { + b.iter_batched( + || (), + |()| start_recursive_control(&fixture.claude), + BatchSize::PerIteration, + ); + }); + + // Tiny tree: both arms should be similar (fixed overhead check). + let tiny = build_fixture(0); + group.throughput(Throughput::Elements(1)); + group.bench_function(BenchmarkId::new("scoped_tiny", 0), |b| { + b.iter_batched(|| (), |()| start_scoped(&tiny), BatchSize::PerIteration); + }); + group.bench_function(BenchmarkId::new("recursive_control_tiny", 0), |b| { + b.iter_batched( + || (), + |()| start_recursive_control(&tiny.claude), + BatchSize::PerIteration, + ); + }); + + group.finish(); +} + +criterion_group!(benches, bench_skills_watcher_startup); +criterion_main!(benches); diff --git a/crates/codegen/xai-grok-shell/changelogs/0.2.115.md b/crates/codegen/xai-grok-shell/changelogs/0.2.115.md index 363b1ca..c0e4f03 100644 --- a/crates/codegen/xai-grok-shell/changelogs/0.2.115.md +++ b/crates/codegen/xai-grok-shell/changelogs/0.2.115.md @@ -1,5 +1,9 @@ # 0.2.115 — 2026-07-29 +## Features + +- **Delete sessions from the dashboard and welcome list.** On the dashboard, press `Ctrl+X` twice (or hover a settled row and click `[✗]` twice); in the welcome and `/resume` lists, press `d` then `y`. + ## Bug Fixes - **Fixed chat history corruption** that could duplicate tool results or cause later 400 errors after repeated identical tool calls. diff --git a/crates/codegen/xai-grok-shell/changelogs/0.2.117.json b/crates/codegen/xai-grok-shell/changelogs/0.2.117.json new file mode 100644 index 0000000..21c4308 --- /dev/null +++ b/crates/codegen/xai-grok-shell/changelogs/0.2.117.json @@ -0,0 +1,37 @@ +[ + { + "category": "fixes", + "description": "**Stop command** now terminates all background subagents from prior turns.", + "breaking_change": false + }, + { + "category": "features", + "description": "**GROK_EXTRA_CA_BUNDLE** env var allows adding custom TLS root certificates.", + "breaking_change": false + }, + { + "category": "fixes", + "description": "**kill_task** tool now correctly reports when a task does not exist over ACP connections.", + "breaking_change": false + }, + { + "category": "fixes", + "description": "**get_task_output** no longer waits the full timeout for already-finished tasks over ACP.", + "breaking_change": false + }, + { + "category": "performance", + "description": "**Terminal resize** is much faster on long conversations in fullscreen mode.", + "breaking_change": false + }, + { + "category": "fixes", + "description": "**/usage** command and billing UI are hidden for enterprise auth setups.", + "breaking_change": false + }, + { + "category": "fixes", + "description": "**Plan approval** no longer starts Build when pressing Enter without notes in revise mode.", + "breaking_change": false + } +] diff --git a/crates/codegen/xai-grok-shell/changelogs/0.2.117.md b/crates/codegen/xai-grok-shell/changelogs/0.2.117.md new file mode 100644 index 0000000..b67d5ef --- /dev/null +++ b/crates/codegen/xai-grok-shell/changelogs/0.2.117.md @@ -0,0 +1,18 @@ +# 0.2.117 — 2026-07-30 + +## Features + +- **GROK_EXTRA_CA_BUNDLE** env var allows adding custom TLS root certificates. + +## Bug Fixes + +- **Stop command** now terminates all background subagents from prior turns. +- **kill_task** tool now correctly reports when a task does not exist over ACP connections. +- **get_task_output** no longer waits the full timeout for already-finished tasks over ACP. +- **/usage** command and billing UI are hidden for enterprise auth setups. +- **Plan approval** no longer starts Build when pressing Enter without notes in revise mode. + +## Performance + +- **Terminal resize** is much faster on long conversations in fullscreen mode. + diff --git a/crates/codegen/xai-grok-shell/src/agent/handlers/session.rs b/crates/codegen/xai-grok-shell/src/agent/handlers/session.rs index 063531c..698bc5c 100644 --- a/crates/codegen/xai-grok-shell/src/agent/handlers/session.rs +++ b/crates/codegen/xai-grok-shell/src/agent/handlers/session.rs @@ -267,8 +267,8 @@ async fn handle_session_list( ) -> Result { use crate::session::unified_list; - // Under chat mode `parse_list_req` REPLACES any client-sent `kind` facet - // (never union) so every list surface is conversations-only. + // Under chat mode `parse_list_req` force-rewrites `kind` to conversations + // unless `local-workspace` is compiled in and the client sent chat/build. let req = unified_list::parse_list_req(args.params.get()) .map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?; tracing::debug!( diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/acp_agent.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/acp_agent.rs index 6ca1c39..bc7a555 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/acp_agent.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/acp_agent.rs @@ -1020,6 +1020,25 @@ impl acp::Agent for MvpAgent { .as_ref() .and_then(|m| m.get("modelId").and_then(|v| v.as_str())) .filter(|s| !s.is_empty()); + #[cfg(all(feature = "local-workspace", unix))] + let pending_local_workspace = self + .start_own_local_workspace_if_needed( + &mut session_meta_for_stamp, + cwd.as_path(), + ) + .await?; + #[cfg(all(feature = "local-workspace", not(unix)))] + { + use crate::gateway_bridge::local_workspace_supervisor::parse_local_workspace_intent; + use crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceIntent; + use crate::gateway_bridge::local_workspace_supervisor::SupervisorError; + if matches!( + parse_local_workspace_intent(session_meta_for_stamp.as_ref()), + Some(LocalWorkspaceIntent::Own { .. }) + ) { + return Err(SupervisorError::UnsupportedPlatform.into_acp_error()); + } + } #[allow(unused_variables)] let session_computer_sessions = resolve_session_computer_sessions( arguments.meta.as_ref(), @@ -1052,6 +1071,15 @@ impl acp::Agent for MvpAgent { } None => acp::SessionId::new(uuid::Uuid::now_v7().to_string()), }; + #[cfg(all(feature = "local-workspace", unix))] + let mut local_ws_reap_guard = self + .new_local_workspace_reap_guard(session_id.clone(), false); + #[cfg(all(feature = "local-workspace", unix))] + if let Some(handle) = pending_local_workspace { + self.register_local_workspace_supervisor(session_id.clone(), handle); + local_ws_reap_guard = self + .new_local_workspace_reap_guard(session_id.clone(), true); + } let mut session_timer = crate::instrumentation_timer!("session.new_session"); session_timer.with_field("session_id", session_id.0.as_ref()); session_timer.with_field("cwd", cwd.as_str()); @@ -1273,8 +1301,16 @@ impl acp::Agent for MvpAgent { }; self.spawn_and_register_session(init, spawn_opts).await }; + #[cfg(all(feature = "local-workspace", unix))] + if spawn_res.is_err() { + self.shutdown_gateway_bridge(&session_id); + } spawn_res?; tracing::debug!(session_id = %session_id.0, "new_session: spawn_session_actor"); + #[cfg(feature = "local-workspace")] + if local_workspace_intent_present(arguments.meta.as_ref()) { + self.mark_local_workspace_bound(session_id.clone()); + } self.maybe_spawn_interactive_trust_prompt( &session_id, cwd.as_path(), @@ -1409,6 +1445,7 @@ impl acp::Agent for MvpAgent { ); insert_applied_tool_overrides(obj, applied_tool_overrides.as_ref()); } + #[cfg(all(feature = "local-workspace", unix))] local_ws_reap_guard.disarm(); Ok( acp::NewSessionResponse::new(session_id) .models(Some(models)) @@ -1947,7 +1984,7 @@ impl acp::Agent for MvpAgent { let persisted_model = summary.current_model_id.clone(); let models = self.models_manager.models(); let available = self.models_manager.available(); - self.model_unavailable_sessions.borrow_mut().remove(session_id.0.as_ref()); + self.session_registry.take_unavailable_model(&session_id); let resolved_catalog_key = resolve_catalog_key(&models, &persisted_model); tracing::debug!( session_id = %session_id.0, @@ -2061,9 +2098,8 @@ impl acp::Agent for MvpAgent { &reason, ) .await; - self.model_unavailable_sessions - .borrow_mut() - .insert(session_id.0.to_string(), persisted_model.clone()); + self.session_registry + .set_unavailable_model(&session_id, persisted_model.clone()); fallback }; tracing::debug!( @@ -2266,10 +2302,8 @@ impl acp::Agent for MvpAgent { return Ok(acp::PromptResponse::new(acp::StopReason::EndTurn)); } let latched_model = self - .model_unavailable_sessions - .borrow() - .get(arguments.session_id.0.as_ref()) - .cloned(); + .session_registry + .unavailable_model(&arguments.session_id); if let Some(unavailable_model) = latched_model { let models = self.models_manager.models(); let available = self.models_manager.available(); @@ -2294,9 +2328,7 @@ impl acp::Agent for MvpAgent { }), ), ); - self.model_unavailable_sessions - .borrow_mut() - .remove(arguments.session_id.0.as_ref()); + self.session_registry.take_unavailable_model(&arguments.session_id); if let Err(e) = crate::agent::handlers::model_switch::apply( self, acp::SetSessionModelRequest::new( @@ -3426,9 +3458,8 @@ impl acp::Agent for MvpAgent { let res = crate::agent::handlers::model_switch::apply(self, args).await; if res.is_ok() && let Some(unavailable) = self - .model_unavailable_sessions - .borrow_mut() - .remove(session_id.0.as_ref()) + .session_registry + .take_unavailable_model(&session_id) { tracing::info!( session_id = %session_id.0, @@ -3488,6 +3519,10 @@ impl acp::Agent for MvpAgent { let ops = self.resolve_workspace_ops()?; crate::extensions::worktree::handle(self, &ops, &args).await } + #[cfg(feature = "local-workspace")] + "x.ai/session/add_local_workspace" => { + crate::extensions::session_admin::handle(self, &args).await + } "x.ai/session/rename" | "x.ai/session/delete" | "x.ai/session/update_mcp_servers" | "x.ai/session/fork" | "x.ai/internal/reload_all_mcp_servers" diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs index 9e90805..9c4ebc3 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs @@ -739,7 +739,7 @@ impl MvpAgent { /// Most recently allocated turn number for `sid`, or `None` if the /// session has not started a turn yet. pub(crate) fn session_turn_number(&self, sid: &acp::SessionId) -> Option { - self.retained_resources.borrow().get(sid).and_then(|d| d.turn_number) + self.session_registry.turn_number(sid) } /// Return the current GrokAuth credentials, if authenticated and not expired. pub(crate) fn current_auth(&self) -> Option { @@ -763,6 +763,596 @@ impl MvpAgent { pub(crate) fn alpha_test_key(&self) -> Option { self.cfg.borrow().endpoints.alpha_test_key.clone() } + #[cfg(all(feature = "local-workspace", unix))] + /// Spawn owned `workspace_server` for chat+local `own` intent. + /// Mints `server_id` into `_meta` before handshake parse. + pub(crate) async fn start_own_local_workspace_if_needed( + &self, + meta: &mut Option, + session_cwd: &std::path::Path, + ) -> Result< + Option, + acp::Error, + > { + use crate::gateway_bridge::local_workspace_supervisor::{ + parse_local_workspace_intent, stamp_server_id_into_meta, start_own, + StartOwnConfig, LocalWorkspaceIntent, + }; + let Some(LocalWorkspaceIntent::Own { cwd }) = parse_local_workspace_intent( + meta.as_ref(), + ) else { + return Ok(None); + }; + let cwd = if cwd.as_os_str().is_empty() { + session_cwd.to_path_buf() + } else { + cwd + }; + crate::gateway_bridge::local_workspace_supervisor::validate_cwd(&cwd) + .map_err(|e| e.into_acp_error())?; + let hub_url = { + let cfg = self.cfg.borrow(); + crate::gateway_bridge::local_workspace_supervisor::resolve_hub_url( + cfg.hub.url.as_deref(), + ) + }; + let handle = start_own(StartOwnConfig { + cwd, + hub_url, + auth_config: None, + binary: None, + ready_timeout: crate::gateway_bridge::local_workspace_supervisor::READY_TIMEOUT, + allow_missing_auth: false, + }) + .await + .map_err(|e| e.into_acp_error())?; + let meta_map = meta.get_or_insert_with(acp::Meta::new); + stamp_server_id_into_meta(meta_map, &handle.server_id); + Ok(Some(handle)) + } + #[cfg(all(feature = "local-workspace", unix))] + pub(crate) fn register_local_workspace_supervisor( + &self, + session_id: acp::SessionId, + handle: crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceHandle, + ) { + let server_id = handle.server_id.clone(); + self.arm_local_workspace_watcher(session_id.clone(), handle); + tracing::info!( + session_id = %session_id.0, + server_id = %server_id, + "local_workspace_supervisor: registered own workspace_server" + ); + } + #[cfg(all(feature = "local-workspace", unix))] + pub(crate) fn new_local_workspace_reap_guard( + &self, + session_id: acp::SessionId, + armed: bool, + ) -> LocalWorkspaceReapGuard { + LocalWorkspaceReapGuard { + supervisors: self.local_workspace_supervisors.clone(), + generations: self.local_workspace_generations.clone(), + session_id, + armed, + } + } + /// Prefer live supervisor `server_id` over the parse-time stamp (pre-bridge crash). + #[cfg(all(feature = "local-workspace", unix))] + pub(crate) fn refresh_sessions_from_supervisor( + &self, + session_id: &acp::SessionId, + sessions: Option>, + ) -> Option> { + use crate::gateway_bridge::ComputerSession; + let supervisors = self.local_workspace_supervisors.borrow(); + let Some(handle) = supervisors.get(session_id) else { + return sessions; + }; + let server_id = handle.server_id.clone(); + let cwd = Some(handle.cwd.to_string_lossy().into_owned()); + match sessions { + None => Some(vec![ComputerSession::ExistingWorkspace { server_id, cwd }]), + Some(mut list) => { + for session in &mut list { + if let ComputerSession::ExistingWorkspace { + server_id: sid, + cwd: existing_cwd, + } = session { + *sid = server_id.clone(); + if existing_cwd.is_none() { + *existing_cwd = cwd.clone(); + } + } + } + Some(list) + } + } + } + /// Wait out an in-flight crash restart before refreshing handshake sessions. + #[cfg(all(feature = "local-workspace", unix))] + pub(crate) async fn await_refresh_sessions_from_supervisor( + &self, + session_id: &acp::SessionId, + sessions: Option>, + ) -> Option> { + const WAIT: std::time::Duration = std::time::Duration::from_secs(5); + let deadline = tokio::time::Instant::now() + WAIT; + loop { + let pending = self + .local_workspace_restart_pending + .borrow() + .contains(session_id); + let live = self + .local_workspace_supervisors + .borrow() + .contains_key(session_id); + if live || !pending { + break; + } + if tokio::time::Instant::now() >= deadline { + tracing::warn!( + session_id = %session_id.0, + "timed out waiting for local-workspace crash restart before handshake refresh" + ); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + self.refresh_sessions_from_supervisor(session_id, sessions) + } + #[cfg(all(feature = "local-workspace", unix))] + fn arm_local_workspace_watcher( + &self, + session_id: acp::SessionId, + handle: crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceHandle, + ) { + let cwd = handle.cwd.clone(); + let sessions = self.session_registry.clone(); + let supervisors = self.local_workspace_supervisors.clone(); + let generations = self.local_workspace_generations.clone(); + let sid = session_id.clone(); + let hub_url = { + let cfg = self.cfg.borrow(); + crate::gateway_bridge::local_workspace_supervisor::resolve_hub_url( + cfg.hub.url.as_deref(), + ) + }; + let auth_path = xai_grok_workspace::hub_auth::default_auth_path().ok(); + let binary = crate::gateway_bridge::local_workspace_supervisor::resolve_workspace_server_bin() + .ok(); + let agent_ref = LocalRef::new(self); + let generation = { + let mut gens = generations.borrow_mut(); + let e = gens.entry(session_id.clone()).or_insert(0); + *e = e.saturating_add(1); + *e + }; + self.local_workspace_supervisors.borrow_mut().insert(session_id.clone(), handle); + let mut supervisors_mut = self.local_workspace_supervisors.borrow_mut(); + let Some(handle_mut) = supervisors_mut.get_mut(&session_id) else { + return; + }; + let _ = handle_mut + .spawn_exit_watcher(move || { + let sessions = sessions.clone(); + let supervisors = supervisors.clone(); + let generations = generations.clone(); + let sid = sid.clone(); + let hub_url = hub_url.clone(); + let auth_path = auth_path.clone(); + let binary = binary.clone(); + let cwd = cwd.clone(); + let agent_ref = agent_ref.clone(); + tokio::task::spawn_local(async move { + if generations.borrow().get(&sid) != Some(&generation) { + return; + } + agent_ref + .get() + .local_workspace_restart_pending + .borrow_mut() + .insert(sid.clone()); + let prev = supervisors.borrow_mut().remove(&sid); + let Some(prev) = prev else { + agent_ref + .get() + .local_workspace_restart_pending + .borrow_mut() + .remove(&sid); + return; + }; + let Some(binary) = binary else { + tracing::warn!( + session_id = %sid.0, + "local workspace crash restart skipped: binary missing" + ); + prev.shutdown().await; + agent_ref + .get() + .local_workspace_restart_pending + .borrow_mut() + .remove(&sid); + return; + }; + let auth = auth_path + .unwrap_or_else(|| std::path::PathBuf::from("/nonexistent")); + let restart_count = prev.restart_count; + let prev_cwd = prev.cwd.clone(); + prev.shutdown().await; + match crate::gateway_bridge::local_workspace_supervisor::restart_own_from( + restart_count, + prev_cwd, + &binary, + &hub_url, + &auth, + false, + ) + .await + { + Ok(new_handle) => { + if generations.borrow().get(&sid) != Some(&generation) { + new_handle.shutdown().await; + agent_ref + .get() + .local_workspace_restart_pending + .borrow_mut() + .remove(&sid); + return; + } + let new_id = new_handle.server_id.clone(); + let cwd_str = cwd.to_string_lossy().into_owned(); + agent_ref + .get() + .arm_local_workspace_watcher(sid.clone(), new_handle); + agent_ref + .get() + .local_workspace_restart_pending + .borrow_mut() + .remove(&sid); + let armed_generation = generations + .borrow() + .get(&sid) + .copied(); + let bridge = sessions.bridge(&sid); + if let Some(bridge) = bridge { + let _ = crate::gateway_bridge::local_workspace_supervisor::push_computer_sessions_update( + &bridge, + new_id.clone(), + Some(cwd_str.clone()), + false, + ) + .await; + let _ = bridge.wait_until_ready().await; + if generations.borrow().get(&sid) + != armed_generation.as_ref() + { + tracing::debug!( + session_id = %sid.0, + "skip stale local-workspace session.update after superseded restart" + ); + return; + } + let _ = crate::gateway_bridge::local_workspace_supervisor::push_computer_sessions_update( + &bridge, + new_id, + Some(cwd_str), + false, + ) + .await; + } + } + Err(err) => { + tracing::warn!( + session_id = %sid.0, + error = %err, + "local workspace crash restart failed" + ); + agent_ref + .get() + .local_workspace_restart_pending + .borrow_mut() + .remove(&sid); + } + } + }); + }); + } + /// add-only mid-session local workspace via ACP extension / session.update. + /// + /// Refuses if a local existing workspace is already bound (no remove until session end). + /// Own mode requires unix (supervisor spawn). Attach is platform-agnostic. + #[cfg(feature = "local-workspace")] + pub(crate) async fn add_local_workspace_mid_session( + &self, + session_id: &acp::SessionId, + mut meta: Option, + session_cwd: &std::path::Path, + ) -> Result { + use crate::gateway_bridge::ComputerSession; + use crate::gateway_bridge::local_workspace_supervisor::{ + parse_local_workspace_intent, LocalWorkspaceIntent, SupervisorError, + }; + if self.local_workspace_already_bound(session_id) { + return Err( + acp::Error::invalid_params() + .data( + serde_json::json!({ + "code": "local_workspace_already_bound", + "message": "local workspace already bound; remove is not supported until session end", + }), + ), + ); + } + self.mark_local_workspace_bound(session_id.clone()); + let mut bind_guard = LocalWorkspaceBindGuard { + bound: self.local_workspace_bound.clone(), + session_id: session_id.clone(), + keep: false, + }; + let Some(intent) = parse_local_workspace_intent(meta.as_ref()) else { + return Err( + acp::Error::invalid_params() + .data( + serde_json::json!({ + "code": "local_workspace_intent_missing", + "message": "x.ai/local_workspace intent required for mid-session add", + }), + ), + ); + }; + let mode = match &intent { + LocalWorkspaceIntent::Own { .. } => "own", + LocalWorkspaceIntent::Attach { .. } => "attach", + }; + #[cfg_attr(not(unix), allow(unused_variables))] + let pending: Option< + crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceHandle, + > = match intent { + LocalWorkspaceIntent::Own { .. } => { + #[cfg(unix)] + { + self.start_own_local_workspace_if_needed(&mut meta, session_cwd) + .await? + } + #[cfg(not(unix))] + { + let _ = session_cwd; + return Err(SupervisorError::UnsupportedPlatform.into_acp_error()); + } + } + LocalWorkspaceIntent::Attach { cwd, .. } => { + if let Some(ref cwd) = cwd { + crate::gateway_bridge::local_workspace_supervisor::validate_cwd(cwd) + .map_err(|e| e.into_acp_error())?; + } + Self::ensure_attach_fs_only_advertised_tools() + .map_err(|msg| { + acp::Error::invalid_params() + .data( + serde_json::json!({ + "code": "local_workspace_fs_only_required", + "message": msg, + }), + ) + })?; + None + } + }; + let sessions = resolve_session_computer_sessions(meta.as_ref())?; + let Some(sessions) = sessions.filter(|s| !s.is_empty()) else { + return Err( + acp::Error::invalid_params() + .data( + serde_json::json!({ + "code": "local_workspace_stamp_failed", + "message": "failed to resolve existing_workspace stamp for mid-session add", + }), + ), + ); + }; + if !sessions + .iter() + .any(|s| matches!(s, ComputerSession::ExistingWorkspace { .. })) + { + return Err( + acp::Error::invalid_params() + .data( + serde_json::json!({ + "code": "local_workspace_stamp_failed", + "message": "mid-session add did not produce existing_workspace", + }), + ), + ); + } + #[cfg(unix)] + let mut reap_guard = self + .new_local_workspace_reap_guard(session_id.clone(), false); + #[cfg(unix)] + if let Some(handle) = pending { + self.register_local_workspace_supervisor(session_id.clone(), handle); + reap_guard = self.new_local_workspace_reap_guard(session_id.clone(), true); + } + let Some(bridge) = self.gateway_bridge_for(session_id) else { + return Err( + acp::Error::invalid_params() + .data( + serde_json::json!({ + "code": "gateway_bridge_missing", + "message": "session has no gateway bridge for session.update computer_sessions", + }), + ), + ); + }; + match tokio::time::timeout(BRIDGE_READY_TIMEOUT, bridge.wait_until_ready()).await + { + Ok(Ok(())) => {} + Ok(Err(err)) => { + return Err(err.into_acp_error()); + } + Err(_) => { + return Err( + acp::Error::internal_error() + .data( + "gateway bridge not ready for mid-session add_local_workspace", + ), + ); + } + } + let sessions = sessions; + let server_id = match sessions.first() { + Some(ComputerSession::ExistingWorkspace { server_id, .. }) => { + server_id.clone() + } + _ => { + return Err( + acp::Error::internal_error() + .data("expected existing_workspace as first computer session"), + ); + } + }; + let cwd = match sessions.first() { + Some(ComputerSession::ExistingWorkspace { cwd, .. }) => cwd.clone(), + _ => None, + }; + if let Some(ref stamped_cwd) = cwd { + crate::gateway_bridge::local_workspace_supervisor::validate_cwd( + std::path::Path::new(stamped_cwd), + ) + .map_err(|e| e.into_acp_error())?; + } + if let Err(err) = crate::gateway_bridge::local_workspace_supervisor::push_computer_sessions_update( + &bridge, + server_id.clone(), + cwd, + true, + ) + .await + { + return Err(err.into_acp_error()); + } + #[cfg(unix)] reap_guard.disarm(); + bind_guard.keep = true; + Ok(serde_json::json!({ + "ok": true, + "server_id": server_id, + "mode": mode, + })) + } + #[cfg(feature = "local-workspace")] + pub(crate) fn local_workspace_already_bound( + &self, + session_id: &acp::SessionId, + ) -> bool { + if self.local_workspace_bound.borrow().contains(session_id) { + return true; + } + #[cfg(unix)] + if self.local_workspace_supervisors.borrow().contains_key(session_id) { + return true; + } + false + } + #[cfg(feature = "local-workspace")] + pub(crate) fn mark_local_workspace_bound(&self, session_id: acp::SessionId) { + self.local_workspace_bound.borrow_mut().insert(session_id); + } + /// Operator-attested FS-only toolset for mid-session attach. + #[cfg(feature = "local-workspace")] + fn ensure_attach_fs_only_advertised_tools() -> Result<(), String> { + const ENV: &str = "GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS"; + const ALLOW: &[&str] = &[ + "workspace.fs_list", + "workspace.fs_exists", + "workspace.fs_read_file", + "workspace.fs_write_file", + "workspace.fs_delete_file", + "workspace.put_files", + "workspace.get_files", + ]; + let Some(raw) = std::env::var(ENV) + .ok() + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) else { + return Err( + "attached workspace_server advertised toolset is uncheckable; refuse attach \ + (set GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS to a comma-separated FS-only catalog)" + .into(), + ); + }; + let ids: Vec<&str> = raw + .split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .collect(); + if ids.is_empty() { + return Err( + "attached workspace_server advertised an empty toolset; refuse attach" + .into(), + ); + } + let forbidden: Vec<&str> = ids + .into_iter() + .filter(|id| !ALLOW.contains(id)) + .collect(); + if forbidden.is_empty() { + Ok(()) + } else { + Err( + format!( + "attached workspace_server advertises tools outside the FS-only allowlist: {}", + forbidden.join(", ") + ), + ) + } + } + #[cfg(feature = "local-workspace")] + /// After chat+local stamp, wait for handshake success. + /// + /// Only fail-closed for `x.ai/local_workspace` intent (not generic + /// GatewayAttach). Handshake errors propagate; session + bridge are reaped + /// on failure / timeout. + pub(crate) async fn await_existing_workspace_handshake( + &self, + session_id: &acp::SessionId, + local_workspace_intent: bool, + ) -> Result<(), acp::Error> { + if !local_workspace_intent { + return Ok(()); + } + #[cfg(feature = "local-workspace")] + self.mark_local_workspace_bound(session_id.clone()); + let Some(bridge) = self.gateway_bridge_for(session_id) else { + return Ok(()); + }; + match tokio::time::timeout(BRIDGE_READY_TIMEOUT, bridge.wait_until_ready()).await + { + Ok(Ok(())) => Ok(()), + Ok(Err(err)) => { + tracing::warn!( + session_id = %session_id.0, + error = %err, + kind = "existing_workspace_handshake_failed", + "chat+local handshake failed; reaping session" + ); + self.request_session_shutdown(session_id); + self.remove_session(session_id); + Err(err.into_acp_error()) + } + Err(_) => { + tracing::warn!( + session_id = %session_id.0, + kind = "existing_workspace_handshake_timeout", + "chat+local handshake timed out; reaping session" + ); + self.request_session_shutdown(session_id); + self.remove_session(session_id); + Err( + acp::Error::internal_error().data("gateway bridge connect timed out"), + ) + } + } + } /// Build the process-lifetime local `WorkspaceOps` on first use. /// /// Deferred past ACP wiring so `initialize` can respond before folder-trust @@ -1896,9 +2486,8 @@ impl MvpAgent { let instance = Self { sessions: RefCell::new(HashMap::new()), activity, + session_registry: SessionRegistry::default(), loading_sessions: RefCell::new(HashMap::new()), - retained_resources: RefCell::new(HashMap::new()), - session_threads: RefCell::new(HashMap::new()), resident_roster_titles: RefCell::new(HashMap::new()), initialize_request: OnceLock::new(), gateway, @@ -1948,7 +2537,6 @@ impl MvpAgent { codebase_indexes: Arc::new( parking_lot::Mutex::new(CodebaseIndexManager::new()), ), - resident_resources: RefCell::new(HashMap::new()), worktree_type, restore_code, session_registry_local, @@ -1958,7 +2546,6 @@ impl MvpAgent { crate::session::mcp_servers::McpState::new(vec![]), ), ), - model_unavailable_sessions: RefCell::new(std::collections::HashMap::new()), subagent_event_tx, subagent_event_rx: RefCell::new(Some(subagent_event_rx)), subagent_presentation: RefCell::new( @@ -1970,7 +2557,18 @@ impl MvpAgent { std::sync::atomic::AtomicBool::new(false), ), workspace_ops: RefCell::new(None), - session_live_state: RefCell::new(HashMap::new()), + #[cfg(all(feature = "local-workspace", unix))] + local_workspace_supervisors: Rc::new(RefCell::new(HashMap::new())), + #[cfg(all(feature = "local-workspace", unix))] + local_workspace_generations: Rc::new(RefCell::new(HashMap::new())), + #[cfg(all(feature = "local-workspace", unix))] + local_workspace_restart_pending: Rc::new( + RefCell::new(std::collections::HashSet::new()), + ), + #[cfg(feature = "local-workspace")] + local_workspace_bound: Rc::new( + RefCell::new(std::collections::HashSet::new()), + ), supervisor_started: std::cell::Cell::new(false), settings_reapply_in_flight: std::rc::Rc::new(std::cell::Cell::new(false)), post_auth_settings_in_flight: std::rc::Rc::new(std::cell::Cell::new(false)), @@ -2074,7 +2672,7 @@ impl MvpAgent { } self.request_session_shutdown(&id); if self.take_session(&id).is_some() { - self.resident_resources.borrow_mut().remove(&id); + self.session_registry.clear_resident(&id); self.set_session_live_state(&id, SessionLiveState::Dormant); unloaded += 1; tracing::debug!(session_id = %id.0, "idle session unloaded to disk on disconnect"); @@ -2095,10 +2693,13 @@ impl MvpAgent { /// Uses async polling (never blocks the `LocalSet` runtime) with a 5s deadline /// to handle slow shutdowns (e.g., embedding API timeouts). pub(super) async fn drain_old_session_thread(&self, session_id: &acp::SessionId) { - let thread = self.session_threads.borrow_mut().remove(session_id); - let Some(thread) = thread else { return }; - if thread.is_finished() { - return; + match self.session_registry.thread_is_finished(session_id) { + None => return, + Some(true) => { + self.session_registry.clear_thread(session_id); + return; + } + Some(false) => {} } tracing::info!( session_id = %session_id.0, @@ -2106,12 +2707,17 @@ impl MvpAgent { ); let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5); loop { - if thread.is_finished() { - tracing::debug!( - session_id = %session_id.0, - "Old session thread finished cleanly" - ); - return; + match self.session_registry.thread_is_finished(session_id) { + None => return, + Some(true) => { + self.session_registry.clear_thread(session_id); + tracing::debug!( + session_id = %session_id.0, + "Old session thread finished cleanly" + ); + return; + } + Some(false) => {} } if tokio::time::Instant::now() >= deadline { tracing::warn!( @@ -2904,12 +3510,7 @@ impl MvpAgent { } /// Set a session's next trace turn number. pub(super) fn set_turn_number(&self, session_id: &acp::SessionId, next: u64) { - self - .retained_resources - .borrow_mut() - .entry(session_id.clone()) - .or_default() - .turn_number = Some(next); + self.session_registry.set_turn_number(session_id, next); } /// Upload each drained harness trace turn as its own `turn_{N}` artifact, /// numbered from the same counter as model turns so subagents interleave @@ -4102,9 +4703,7 @@ impl MvpAgent { ) .await? }; - self.session_threads - .borrow_mut() - .insert(session_info.id.clone(), session_thread); + self.session_registry.set_thread(&session_info.id, session_thread); tracing::debug!(session_id = %session_info.id.0, "spawn_session_on_thread complete"); self.set_session_live_state(&session_info.id, SessionLiveState::IdleResident); self.ensure_session_supervisor(); @@ -4164,12 +4763,8 @@ impl MvpAgent { } }); } - self - .retained_resources - .borrow_mut() - .entry(session_info.id.clone()) - .or_default() - .permission_event_receiver = Some(permission_events_rx); + self.session_registry + .set_permission_receiver(&session_info.id, permission_events_rx); if handle_display_cwd.is_some() { handle.display_cwd = handle_display_cwd; } @@ -4203,16 +4798,56 @@ impl MvpAgent { &self, session_id: &acp::SessionId, ) -> Vec { - let mut events = Vec::new(); - let mut retained = self.retained_resources.borrow_mut(); - if let Some(rx) = retained - .get_mut(session_id) - .and_then(|d| d.permission_event_receiver.as_mut()) - { - while let Ok(event) = rx.try_recv() { - events.push(event); - } - } - events + self.session_registry.drain_permission_events(session_id) + } +} +/// Rollback guard for mid-session bind reservation. +#[cfg(feature = "local-workspace")] +struct LocalWorkspaceBindGuard { + bound: Rc>>, + session_id: acp::SessionId, + keep: bool, +} +#[cfg(feature = "local-workspace")] +impl Drop for LocalWorkspaceBindGuard { + fn drop(&mut self) { + if !self.keep { + self.bound.borrow_mut().remove(&self.session_id); + } + } +} +/// Reap guard: if session/new fails after register, Drop kills the supervisor. +#[cfg(all(feature = "local-workspace", unix))] +pub(crate) struct LocalWorkspaceReapGuard { + supervisors: Rc< + RefCell< + HashMap< + acp::SessionId, + crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceHandle, + >, + >, + >, + generations: Rc>>, + session_id: acp::SessionId, + armed: bool, +} +#[cfg(all(feature = "local-workspace", unix))] +impl LocalWorkspaceReapGuard { + pub(crate) fn disarm(&mut self) { + self.armed = false; + } +} +#[cfg(all(feature = "local-workspace", unix))] +impl Drop for LocalWorkspaceReapGuard { + fn drop(&mut self) { + if !self.armed { + return; + } + self.generations.borrow_mut().remove(&self.session_id); + if let Some(handle) = self.supervisors.borrow_mut().remove(&self.session_id) { + tokio::spawn(async move { + handle.shutdown().await; + }); + } } } diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/code_nav.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/code_nav.rs index 6e675c8..94e73b2 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/code_nav.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/code_nav.rs @@ -34,11 +34,8 @@ impl MvpAgent { // Pin the index to the requesting session so the Weak in // CodebaseIndexManager doesn't orphan it immediately. if let Some(sid) = session_id { - self.resident_resources - .borrow_mut() - .entry(sid.clone()) - .or_default() - .codebase_index = Some(std::sync::Arc::clone(&handle)); + self.session_registry + .set_codebase_index(sid, std::sync::Arc::clone(&handle)); } Some((handle, was_newly_started)) } diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/mod.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/mod.rs index 34485b5..0a9bf85 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/mod.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/mod.rs @@ -183,6 +183,51 @@ pub(crate) fn jwt_claim_matches_user_subscription_tier( _ => jwt_claim.parse::().is_ok_and(|n| n != 0), } } +/// ACP `_meta` key for chat+local workspace intent (pager stamps on chat create). +#[cfg(feature = "local-workspace")] +const LOCAL_WORKSPACE_META_KEY: &str = "x.ai/local_workspace"; +/// True when `_meta` carries a valid chat+local intent object +/// (`mode` is `"own"` or `"attach"`). +#[cfg(feature = "local-workspace")] +fn local_workspace_intent_present(meta: Option<&acp::Meta>) -> bool { + meta.and_then(|m| m.get(LOCAL_WORKSPACE_META_KEY)) + .and_then(|v| v.as_object()) + .and_then(|o| o.get("mode")) + .and_then(|m| m.as_str()) + .is_some_and(|mode| mode == "own" || mode == "attach") +} +/// valid local-workspace intent → ExistingWorkspace only. +/// +/// `server_id` comes from the intent object, else `cloud_existing_workspace`. +/// Never reads `envId` / never emits `SandboxEnvironment`. +#[cfg(feature = "local-workspace")] +fn parse_local_workspace_existing( + meta: Option<&acp::Meta>, +) -> Option { + use crate::gateway_bridge::ComputerSession; + let local = meta.and_then(|m| m.get(LOCAL_WORKSPACE_META_KEY))?; + let mode = local.get("mode").and_then(|v| v.as_str())?; + if mode != "own" && mode != "attach" { + return None; + } + let server_id = meta_non_empty_str(local, "server_id") + .or_else(|| { + meta + .and_then(|m| m.get(CLOUD_EXISTING_WORKSPACE_META_KEY)) + .and_then(|w| meta_non_empty_str(w, "server_id")) + })?; + let cwd = meta_non_empty_str(local, "cwd") + .or_else(|| { + meta + .and_then(|m| m.get(CLOUD_EXISTING_WORKSPACE_META_KEY)) + .and_then(|w| meta_non_empty_str(w, "cwd")) + }); + Some(ComputerSession::ExistingWorkspace { + server_id, + cwd, + }) +} +#[allow(dead_code)] fn parse_session_computer_sessions(_meta: Option<&acp::Meta>) -> Option> { None } @@ -666,16 +711,12 @@ pub struct MvpAgent { /// leader's auto-update checker, which cannot read the `!Send` maps. Expires /// when the actor exits. See [`crate::agent::activity::AgentActivity`]. pub(crate) activity: crate::agent::activity::AgentActivity, - /// LEADER-SAFE(per-session): in-flight `session/load` guards. Lets a racing - /// `session/prompt` wait via [`Self::wait_for_in_flight_session_load`] instead - /// of failing "unknown session id"; the RAII guard's drop wakes waiters. + /// LEADER-SAFE(per-session). + session_registry: SessionRegistry, + /// A load guard rather than session state: it exists before the session. loading_sessions: RefCell< HashMap>, >, - /// LEADER-SAFE(per-session): reclaimed at `remove_session`. See [`RetainedResources`]. - retained_resources: RefCell>, - /// LEADER-SAFE(per-session): keyed by SessionId. Mirrors `sessions` lifecycle. - session_threads: RefCell>, /// Title per resident session id, refreshed each `build_roster`. Lets the /// synchronous roster deltas reuse the title instead of emitting an empty /// one — `resident_roster_entry` can't read disk. @@ -793,8 +834,6 @@ pub struct MvpAgent { /// LEADER-SAFE(shared): agent-level code-nav index manager, keyed by cwd, /// no per-client state. codebase_indexes: Arc>, - /// LEADER-SAFE(per-session): reclaimed on removal / idle-unload. See [`ResidentResources`]. - resident_resources: RefCell>, /// Worktree creation type (resolved: local config > remote > default Linked). pub(crate) worktree_type: crate::util::config::WorktreeType, /// Restore codebase state on worktree resume (resolved: local config > remote > default false). @@ -810,15 +849,6 @@ pub struct MvpAgent { agent_mcp_state: std::sync::Arc< tokio::sync::Mutex, >, - /// Sessions whose persisted model was unavailable at `session/load` time - /// with no same-family fallback, keyed by session id → the unavailable - /// model id. Prompts to these sessions are blocked until either - /// (a) the model reappears in the catalog — the catalog can be - /// transiently degraded when a reconnect replays `session/load` (e.g. - /// fetch still in flight after a leader restart), so the prompt path - /// re-checks and self-heals — or (b) the user explicitly switches - /// models via `set_session_model`. Released by `remove_session`. - model_unavailable_sessions: RefCell>, /// Unified sender for all subagent coordinator events. /// LEADER-SAFE(shared): channel is multi-producer, coordinator drains. subagent_event_tx: tokio::sync::mpsc::UnboundedSender< @@ -888,13 +918,28 @@ pub struct MvpAgent { /// The agent never opens Computer Hub as a harness/client; remote cloud /// sandboxes are gateway-owned (`gateway_bridge` / `computer_sessions`). workspace_ops: RefCell>, - /// Per-session coarse lifecycle state (residency + turn-state). - /// Updated by `spawn_and_register_session` (→ `IdleResident`) and the - /// join-handle supervisor on actor exit (→ `DeadFailed`) / explicit close - /// (→ `Completed`). This is the roster's data source in PR-6; for now it - /// gives the supervisor an observable demotion signal. - /// LEADER-SAFE(per-session): keyed by SessionId. - session_live_state: RefCell>, + /// Per-session owned local `workspace_server` handles (chat+local `own`). + #[cfg(all(feature = "local-workspace", unix))] + local_workspace_supervisors: Rc< + RefCell< + HashMap< + acp::SessionId, + crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceHandle, + >, + >, + >, + /// Invalidates in-flight crash restarts when the session supervisor is reaped. + #[cfg(all(feature = "local-workspace", unix))] + local_workspace_generations: Rc>>, + /// Sessions whose own supervisor is mid crash-restart (map entry temporarily empty). + #[cfg(all(feature = "local-workspace", unix))] + local_workspace_restart_pending: Rc< + RefCell>, + >, + /// Sessions that already have a local existing workspace (own or attach). + /// mid-session add refuses while this is set; cleared on session end. + #[cfg(feature = "local-workspace")] + local_workspace_bound: Rc>>, /// Idempotency guard: the join-handle supervisor task is spawned at most /// once (on the first `spawn_and_register_session`). See /// `ensure_session_supervisor`. @@ -1297,10 +1342,12 @@ impl Drop for SessionLoadGuard<'_> { mod code_nav; mod folder_trust_prompt; mod heap_profile; +mod session_registry; mod session_lifecycle; mod subagent_coordinator; mod agent_ops; mod acp_agent; +use session_registry::SessionRegistry; pub(crate) use session_lifecycle::RegistrySnapshot; pub(super) use super::ext_parsers; /// Emit the `auth.lifecycle` login span with optional user id and error diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/session_lifecycle.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/session_lifecycle.rs index 3e58a74..335f8c0 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/session_lifecycle.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/session_lifecycle.rs @@ -26,12 +26,8 @@ impl MvpAgent { }); let _ = handle.cmd_tx.send(SessionCommand::Shutdown); drop(handle); - let thread = self.session_threads.borrow_mut().remove(id); self.remove_session_terminal(id, SessionLiveState::Completed); - if let Some(thread) = thread { - self.session_threads.borrow_mut().insert(id.clone(), thread); - self.drain_old_session_thread(id).await; - } + self.drain_old_session_thread(id).await; } /// Finalize the cloud session replica (fire-and-forget, "Hook 4"). /// @@ -78,13 +74,7 @@ impl MvpAgent { parent_session_id: id.0.to_string(), }); self.take_session(id); - self.session_threads.borrow_mut().remove(id); - self.resident_resources.borrow_mut().remove(id); - self.retained_resources.borrow_mut().remove(id); - self.model_unavailable_sessions - .borrow_mut() - .remove(id.0.as_ref()); - self.session_live_state.borrow_mut().remove(id); + self.session_registry.release(id); if let Some(ops) = self.workspace_ops.borrow().as_ref() { ops.end_local_session(id.0.as_ref()); } @@ -97,13 +87,7 @@ impl MvpAgent { /// Cancels therefore wait behind an intake preamble: keep preambles lean. /// Bridge cancels take their own path and stay unordered against this lock. pub(super) fn dispatch_lock(&self, id: &acp::SessionId) -> std::rc::Rc> { - self.retained_resources - .borrow_mut() - .entry(id.clone()) - .or_default() - .dispatch_lock - .get_or_insert_with(Default::default) - .clone() + self.session_registry.dispatch_lock(id) } /// Close a session in response to an **explicit** terminal close /// (`x.ai/session/close`). Finalizes the cloud replica (genuine session @@ -114,14 +98,12 @@ impl MvpAgent { } /// Record the coarse lifecycle state for a session. pub(super) fn set_session_live_state(&self, id: &acp::SessionId, state: SessionLiveState) { - self.session_live_state - .borrow_mut() - .insert(id.clone(), state); + self.session_registry.set_live(id, state); } /// Read the recorded lifecycle state for a session (test observability). #[cfg(test)] pub(super) fn session_live_state_for(&self, id: &acp::SessionId) -> Option { - self.session_live_state.borrow().get(id).copied() + self.session_registry.live(id) } /// Roster-delta hook for a terminally removed session. Broadcasts an /// `x.ai/sessions/changed` notification with the session in `removed` so @@ -231,7 +213,7 @@ impl MvpAgent { if turn_running { return RosterActivity::Working; } - match self.session_live_state.borrow().get(id).copied() { + match self.session_registry.live(id) { Some(SessionLiveState::Completed) => RosterActivity::Completed, Some(SessionLiveState::DeadFailed) => RosterActivity::Dead, Some(SessionLiveState::Dormant) => RosterActivity::Dormant, @@ -339,13 +321,7 @@ impl MvpAgent { /// check is required. Runs both opportunistically and from the join-handle /// supervisor (`ensure_session_supervisor`). pub(super) fn sweep_dead_sessions(&self) { - let dead: Vec = self - .session_threads - .borrow() - .iter() - .filter(|(_, t)| t.is_finished()) - .map(|(id, _)| id.clone()) - .collect(); + let dead = self.session_registry.finished_threads(); for id in dead { if self.sessions.borrow().contains_key(&id) { tracing::warn!( @@ -354,8 +330,7 @@ impl MvpAgent { ); self.reap_dead_session(&id); } else { - self.session_threads.borrow_mut().remove(&id); - self.session_live_state.borrow_mut().remove(&id); + self.session_registry.clear_exited_thread(&id); tracing::debug!( session_id = %id.0, "Reaped finished thread for non-resident session (clean exit)" @@ -453,8 +428,9 @@ impl MvpAgent { .await .unwrap_or(true) } - /// Entry counts for every collection [`Self::remove_session`] drains, - /// plus workspace bindings and shared coordinator state. + /// Entry counts `remove_session` and `x.ai/debug/agent` care about, which + /// includes maps outside [`SessionResources`]: the handle map, load guards, + /// rewind snapshots, subagents, and workspace bindings. pub(crate) async fn registry_snapshot(&self) -> RegistrySnapshot { let subagents = xai_grok_tools::implementations::grok_build::task::backend::ChannelBackend::new( @@ -462,52 +438,29 @@ impl MvpAgent { ) .registry_counts() .await; - let (resident_resources, session_index_claims, require_gateway_sessions) = { - let resident = self.resident_resources.borrow(); - ( - resident.len(), - resident - .values() - .filter(|r| r.codebase_index.is_some()) - .count(), - resident.values().filter(|r| r.require_gateway).count(), - ) - }; - let retained = self.retained_resources.borrow(); - let retained_resources = retained.len(); - let dispatch_locks = retained - .values() - .filter(|d| d.dispatch_lock.is_some()) - .count(); - let session_turn_numbers = retained - .values() - .filter(|d| d.turn_number.is_some()) - .count(); - let permission_event_receivers = retained - .values() - .filter(|d| d.permission_event_receiver.is_some()) - .count(); - drop(retained); + let workspace_ops = self.workspace_ops.borrow(); + let workspace = workspace_ops + .as_ref() + .and_then(|ops| ops.workspace_handle()); + let counts = self.session_registry.counts(); RegistrySnapshot { sessions: self.sessions.borrow().len(), - session_threads: self.session_threads.borrow().len(), - resident_resources, - retained_resources, - dispatch_locks, - session_turn_numbers, - permission_event_receivers, - model_unavailable_sessions: self.model_unavailable_sessions.borrow().len(), - session_live_state: self.session_live_state.borrow().len(), - session_index_claims, - require_gateway_sessions, + loading_sessions: self.loading_sessions.borrow().len(), + session_threads: counts.session_threads, + resident_resources: counts.resident_resources, + retained_resources: counts.retained_resources, + dispatch_locks: counts.dispatch_locks, + session_turn_numbers: counts.session_turn_numbers, + permission_event_receivers: counts.permission_event_receivers, + model_unavailable_sessions: counts.model_unavailable_sessions, + session_live_state: counts.session_live_state, + session_index_claims: counts.session_index_claims, + require_gateway_sessions: counts.require_gateway_sessions, subagent_pending: subagents.pending, subagent_active: subagents.active, subagent_completed: subagents.completed, - workspace_bindings: self - .workspace_ops - .borrow() - .as_ref() - .and_then(|ops| ops.workspace_handle().map(|h| h.session_count())), + workspace_bindings: workspace.map(|h| h.session_count()), + workspace_activity_sessions: workspace.map(|h| h.activity_tracker().session_count()), } } } @@ -516,6 +469,7 @@ impl MvpAgent { #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)] pub struct RegistrySnapshot { pub sessions: usize, + pub loading_sessions: usize, pub session_threads: usize, pub resident_resources: usize, pub retained_resources: usize, @@ -530,4 +484,5 @@ pub struct RegistrySnapshot { pub subagent_active: usize, pub subagent_completed: usize, pub workspace_bindings: Option, + pub workspace_activity_sessions: Option, } diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/session_registry.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/session_registry.rs new file mode 100644 index 0000000..089b30b --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/session_registry.rs @@ -0,0 +1,238 @@ +//! Per-session resources and the registry that owns them. Distinct from +//! `agent::session_registry_client`, which talks to the remote registry. +use super::*; +/// The map stays private so every caller goes through a named operation. +#[derive(Clone, Default)] +pub(super) struct SessionRegistry { + sessions: Rc>>, +} +/// The per-session state this registry owns: retained, resident, thread, live, +/// unavailable model, and bridge. Load guards, rewind snapshots, local +/// workspaces, and the handle map are owned elsewhere, so a new field belongs +/// here only if `release` should drop it with the rest. +#[derive(Default)] +struct SessionResources { + retained: Option, + /// Cleared at idle-unload; survives a reload rebuild. + resident: Option, + thread: Option, + live: Option, + unavailable_model: Option, +} +#[derive(Default)] +pub(super) struct SessionCounts { + pub(super) retained_resources: usize, + pub(super) resident_resources: usize, + pub(super) session_threads: usize, + pub(super) session_live_state: usize, + pub(super) model_unavailable_sessions: usize, + pub(super) dispatch_locks: usize, + pub(super) session_turn_numbers: usize, + pub(super) permission_event_receivers: usize, + pub(super) session_index_claims: usize, + pub(super) require_gateway_sessions: usize, +} +impl SessionRegistry { + /// Releases everything a closing session leaves behind, in one drop. + /// + /// A running actor thread stays: dropping its handle would detach it, and + /// nothing would track the memory it holds. The sweep reclaims it later. + pub(super) fn release(&self, id: &acp::SessionId) { + let mut entries = self.sessions.borrow_mut(); + let Some(mut released) = entries.remove(id) else { + return; + }; + let running = released.thread.take().filter(|t| !t.is_finished()); + drop(released); + if running.is_some() { + entries.insert( + id.clone(), + SessionResources { + thread: running, + retained: None, + resident: None, + live: None, + unavailable_model: None, + }, + ); + } + } + pub(super) fn set_thread(&self, id: &acp::SessionId, thread: SessionThread) { + let displaced = self.edit(id, |e| e.thread.replace(thread)); + if displaced.is_some_and(|t| !t.is_finished()) { + tracing::warn!(session_id = %id.0, "session thread displaced while still running"); + } + } + /// Drops the tracked thread. Returns nothing on purpose: handing a + /// `SessionThread` to a caller lets the last handle die in a local, which + /// detaches the thread with no record left for the sweep. + pub(super) fn clear_thread(&self, id: &acp::SessionId) { + self.clear(id, |e| e.thread = None); + } + /// `None` when no thread is tracked for the session. + #[cfg(test)] + pub(super) fn has_thread(&self, id: &acp::SessionId) -> bool { + self.with(id, |e| e.thread.is_some()).unwrap_or(false) + } + pub(super) fn thread_is_finished(&self, id: &acp::SessionId) -> Option { + self.with(id, |e| e.thread.as_ref().map(SessionThread::is_finished)) + .flatten() + } + pub(super) fn finished_threads(&self) -> Vec { + self.sessions + .borrow() + .iter() + .filter(|(_, e)| e.thread.as_ref().is_some_and(SessionThread::is_finished)) + .map(|(id, _)| id.clone()) + .collect() + } + pub(super) fn clear_exited_thread(&self, id: &acp::SessionId) { + self.clear(id, |e| { + e.thread = None; + e.live = None; + }); + } + pub(super) fn set_live(&self, id: &acp::SessionId, state: SessionLiveState) { + self.edit(id, |e| e.live = Some(state)); + } + pub(super) fn live(&self, id: &acp::SessionId) -> Option { + self.with(id, |e| e.live).flatten() + } + pub(super) fn clear_resident(&self, id: &acp::SessionId) { + self.clear(id, |e| e.resident = None); + } + pub(super) fn set_unavailable_model(&self, id: &acp::SessionId, model: acp::ModelId) { + self.edit(id, |e| e.unavailable_model = Some(model)); + } + pub(super) fn unavailable_model(&self, id: &acp::SessionId) -> Option { + self.with(id, |e| e.unavailable_model.clone()).flatten() + } + pub(super) fn take_unavailable_model(&self, id: &acp::SessionId) -> Option { + let model = self + .sessions + .borrow_mut() + .get_mut(id) + .and_then(|e| e.unavailable_model.take()); + self.drop_if_empty(id); + model + } + pub(super) fn turn_number(&self, id: &acp::SessionId) -> Option { + self.with(id, |e| e.retained.as_ref()?.turn_number) + .flatten() + } + pub(super) fn set_turn_number(&self, id: &acp::SessionId, next: u64) { + self.edit(id, |e| { + e.retained.get_or_insert_default().turn_number = Some(next); + }); + } + pub(super) fn dispatch_lock(&self, id: &acp::SessionId) -> Rc> { + self.edit(id, |e| { + e.retained + .get_or_insert_default() + .dispatch_lock + .get_or_insert_with(Default::default) + .clone() + }) + } + pub(super) fn set_permission_receiver( + &self, + id: &acp::SessionId, + rx: tokio::sync::mpsc::UnboundedReceiver, + ) { + self.edit(id, |e| { + e.retained.get_or_insert_default().permission_event_receiver = Some(rx); + }); + } + pub(super) fn drain_permission_events(&self, id: &acp::SessionId) -> Vec { + let mut events = Vec::new(); + let mut entries = self.sessions.borrow_mut(); + if let Some(rx) = entries + .get_mut(id) + .and_then(|e| e.retained.as_mut()) + .and_then(|r| r.permission_event_receiver.as_mut()) + { + while let Ok(event) = rx.try_recv() { + events.push(event); + } + } + events + } + pub(super) fn set_codebase_index( + &self, + id: &acp::SessionId, + index: std::sync::Arc, + ) { + self.edit(id, |e| { + e.resident.get_or_insert_default().codebase_index = Some(index); + }); + } + /// Destructured so a new field has to be counted, or go unmeasured. + pub(super) fn counts(&self) -> SessionCounts { + let mut counts = SessionCounts::default(); + for entry in self.sessions.borrow().values() { + let SessionResources { + retained, + resident, + thread, + live, + unavailable_model, + } = entry; + counts.retained_resources += usize::from(retained.is_some()); + counts.resident_resources += usize::from(resident.is_some()); + counts.session_threads += usize::from(thread.is_some()); + counts.session_live_state += usize::from(live.is_some()); + counts.model_unavailable_sessions += usize::from(unavailable_model.is_some()); + if let Some(retained) = retained { + counts.dispatch_locks += usize::from(retained.dispatch_lock.is_some()); + counts.session_turn_numbers += usize::from(retained.turn_number.is_some()); + counts.permission_event_receivers += + usize::from(retained.permission_event_receiver.is_some()); + } + if let Some(resident) = resident { + counts.session_index_claims += usize::from(resident.codebase_index.is_some()); + counts.require_gateway_sessions += usize::from(resident.require_gateway); + } + } + counts + } + fn with(&self, id: &acp::SessionId, f: impl FnOnce(&SessionResources) -> R) -> Option { + self.sessions.borrow().get(id).map(f) + } + fn edit(&self, id: &acp::SessionId, f: impl FnOnce(&mut SessionResources) -> R) -> R { + f(self.sessions.borrow_mut().entry(id.clone()).or_default()) + } + fn clear(&self, id: &acp::SessionId, f: impl FnOnce(&mut SessionResources)) { + { + let mut entries = self.sessions.borrow_mut(); + let Some(entry) = entries.get_mut(id) else { + return; + }; + f(entry); + } + self.drop_if_empty(id); + } + fn drop_if_empty(&self, id: &acp::SessionId) { + let mut entries = self.sessions.borrow_mut(); + if entries.get(id).is_some_and(SessionResources::is_empty) { + entries.remove(id); + } + } +} +impl SessionResources { + fn is_empty(&self) -> bool { + let Self { + retained, + resident, + thread, + live, + unavailable_model, + } = self; + let chat_vacant = true; + retained.is_none() + && resident.is_none() + && thread.is_none() + && live.is_none() + && unavailable_model.is_none() + && chat_vacant + } +} diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs index 519af15..a26eef8 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs @@ -3110,6 +3110,290 @@ fn chat_new_session_model_state_matrix() { ); } } +/// valid `x.ai/local_workspace` → ExistingWorkspace only. +/// Never reads `envId` / never emits SandboxEnvironment. +#[cfg(feature = "local-workspace")] +#[test] +fn parse_session_computer_sessions_local_workspace_matrix() { + use crate::gateway_bridge::ComputerSession; + use serde_json::json; + fn existing(server_id: &str, cwd: Option<&str>) -> Vec { + vec![ComputerSession::ExistingWorkspace { + server_id: server_id.to_owned(), + cwd: cwd.map(str::to_owned), + }] + } + let cases: &[(&str, serde_json::Value, Option>)] = &[ + ( + "attach_server_id_on_local", + json!({ + "x.ai/local_workspace": { + "mode": "attach", + "server_id": "lw-attach-1", + "cwd": "/repo", + }, + "envId": "env-must-be-ignored", + }), + Some(existing("lw-attach-1", Some("/repo"))), + ), + ( + "attach_server_id_from_cloud_existing", + json!({ + "x.ai/local_workspace": { + "mode": "attach", + "cwd": "/repo", + }, + "x.ai/cloud_existing_workspace": { + "server_id": "lw-attach-2", + "cwd": "/repo-existing", + }, + "envId": "env-must-be-ignored", + }), + Some(existing("lw-attach-2", Some("/repo"))), + ), + ( + "own_with_server_id_ignores_envid", + json!({ + "x.ai/local_workspace": { + "mode": "own", + "server_id": "lw-own-1", + "cwd": "/Users/me/src", + }, + "envId": "env-must-be-ignored", + }), + Some(existing("lw-own-1", Some("/Users/me/src"))), + ), + ( + "own_without_server_id_no_sandbox_fallback", + json!({ + "x.ai/local_workspace": { + "mode": "own", + "cwd": "/Users/me/src", + }, + "envId": "env-must-be-ignored", + }), + None, + ), + ( + "invalid_mode_falls_through_to_envid", + json!({ + "x.ai/local_workspace": { + "mode": "bogus", + "server_id": "lw-x", + }, + "envId": "env-prod", + }), + Some(vec![ComputerSession::SandboxEnvironment { + environment_id: Some("env-prod".to_owned()), + }]), + ), + ( + "non_object_local_falls_through_to_envid", + json!({ + "x.ai/local_workspace": "not-an-object", + "envId": "env-prod", + }), + Some(vec![ComputerSession::SandboxEnvironment { + environment_id: Some("env-prod".to_owned()), + }]), + ), + ]; + for (label, meta, expected) in cases { + let got = parse_session_computer_sessions(meta.as_object()); + assert_eq!( + got.as_deref(), + expected.as_deref(), + "[{label}] local_workspace match-table mismatch" + ); + } +} +/// Local intent without resolvable server_id fails closed (no silent unstamped start). +#[cfg(feature = "local-workspace")] +#[test] +fn resolve_local_workspace_missing_server_id_fails_closed() { + use serde_json::json; + let meta = json!({ + "x.ai/session": { "kind": "chat" }, + "x.ai/local_workspace": { + "mode": "own", + "cwd": "/repo", + } + }); + let err = resolve_session_computer_sessions(meta.as_object()) + .expect_err("own without server_id must fail closed"); + assert_eq!( + err.data + .as_ref() + .and_then(|d| d.get("code")) + .and_then(|v| v.as_str()), + Some("local_workspace_server_id_missing") + ); +} +/// Supervisor map + reap guard / shutdown_gateway_bridge tear down the entry. +#[cfg(all(feature = "local-workspace", unix))] +#[test] +fn local_workspace_reap_guard_and_shutdown_clear_map() { + run_local_for_bridge_test(|| async { + let agent = build_minimal_agent_for_tests(); + let sid = gateway_bridge_test_session_id(); + { + let mut guard = agent.new_local_workspace_reap_guard(sid.clone(), true); + guard.disarm(); + } + assert!(agent.local_workspace_supervisors.borrow().is_empty()); + agent.shutdown_gateway_bridge(&sid); + assert!( + agent + .local_workspace_generations + .borrow() + .get(&sid) + .is_none() + ); + }); +} +/// Pre-bridge crash refresh rewrites handshake stamp from live supervisor id. +#[cfg(all(feature = "local-workspace", unix))] +#[test] +fn refresh_sessions_from_supervisor_overrides_server_id() { + use crate::gateway_bridge::ComputerSession; + use crate::gateway_bridge::local_workspace_supervisor::test_start_ready_own; + run_local_for_bridge_test(|| async { + let agent = build_minimal_agent_for_tests(); + let sid = gateway_bridge_test_session_id(); + let original = Some(vec![ComputerSession::ExistingWorkspace { + server_id: "lw-stale".into(), + cwd: Some("/repo".into()), + }]); + let unchanged = agent.refresh_sessions_from_supervisor(&sid, original.clone()); + assert!(matches!( + unchanged.as_ref().and_then(|v| v.first()), + Some(ComputerSession::ExistingWorkspace { server_id, .. }) if server_id == "lw-stale" + )); + let (_dir, handle) = test_start_ready_own().await; + let live_id = handle.server_id.clone(); + agent.register_local_workspace_supervisor(sid.clone(), handle); + let refreshed = agent.refresh_sessions_from_supervisor(&sid, original); + match refreshed.as_ref().and_then(|v| v.first()) { + Some(ComputerSession::ExistingWorkspace { server_id, .. }) => { + assert_eq!( + server_id, &live_id, + "refresh must use live supervisor server_id" + ); + } + other => panic!("expected ExistingWorkspace, got {other:?}"), + } + agent.shutdown_gateway_bridge(&sid); + }); +} +/// start_own + register stamps server_id into meta and stores the handle. +#[cfg(all(feature = "local-workspace", unix))] +#[test] +fn start_own_registers_and_stamps_server_id() { + use crate::gateway_bridge::local_workspace_supervisor::{ + stamp_server_id_into_meta, test_start_ready_own, + }; + run_local_for_bridge_test(|| async { + let agent = build_minimal_agent_for_tests(); + let sid = gateway_bridge_test_session_id(); + let (_dir, handle) = test_start_ready_own().await; + let server_id = handle.server_id.clone(); + let mut meta = acp::Meta::new(); + meta.insert( + "x.ai/local_workspace".into(), + serde_json::json!({"mode": "own", "cwd": "/tmp/repo"}), + ); + stamp_server_id_into_meta(&mut meta, &server_id); + assert_eq!( + meta.get("x.ai/local_workspace") + .and_then(|v| v.get("server_id")) + .and_then(|v| v.as_str()), + Some(server_id.as_str()) + ); + agent.register_local_workspace_supervisor(sid.clone(), handle); + assert!( + agent + .local_workspace_supervisors + .borrow() + .contains_key(&sid), + "handle must be registered by SessionId" + ); + assert!( + agent + .local_workspace_generations + .borrow() + .get(&sid) + .is_some_and(|g| *g >= 1), + "arm must bump generation" + ); + agent.shutdown_gateway_bridge(&sid); + }); +} +/// Armed reap guard removes a registered supervisor on drop (session/new failure). +#[cfg(all(feature = "local-workspace", unix))] +#[test] +fn reap_guard_drop_removes_registered_supervisor() { + use crate::gateway_bridge::local_workspace_supervisor::test_start_ready_own; + run_local_for_bridge_test(|| async { + let agent = build_minimal_agent_for_tests(); + let sid = gateway_bridge_test_session_id(); + let (_dir, handle) = test_start_ready_own().await; + agent.register_local_workspace_supervisor(sid.clone(), handle); + assert!( + agent + .local_workspace_supervisors + .borrow() + .contains_key(&sid) + ); + { + let _guard = agent.new_local_workspace_reap_guard(sid.clone(), true); + } + assert!( + agent + .local_workspace_supervisors + .borrow() + .get(&sid) + .is_none(), + "armed guard drop must reap supervisor" + ); + assert!( + agent + .local_workspace_generations + .borrow() + .get(&sid) + .is_none(), + "armed guard drop must invalidate generation" + ); + }); +} +/// Shutdown generation invalidates a pending restart re-insert. +#[cfg(all(feature = "local-workspace", unix))] +#[test] +fn shutdown_generation_invalidates_stale_restart() { + use crate::gateway_bridge::local_workspace_supervisor::test_start_ready_own; + run_local_for_bridge_test(|| async { + let agent = build_minimal_agent_for_tests(); + let sid = gateway_bridge_test_session_id(); + let (_dir, handle) = test_start_ready_own().await; + agent.register_local_workspace_supervisor(sid.clone(), handle); + let generation = *agent + .local_workspace_generations + .borrow() + .get(&sid) + .expect("generation after register"); + agent.shutdown_gateway_bridge(&sid); + assert!( + agent.local_workspace_generations.borrow().get(&sid) != Some(&generation), + "shutdown must invalidate generation so stale restart cannot re-insert" + ); + assert!( + agent + .local_workspace_supervisors + .borrow() + .get(&sid) + .is_none() + ); + }); +} /// `spawn_gateway_bridge` uses `tokio::task::spawn_local`. fn run_local_for_bridge_test(body: F) -> T where @@ -3174,39 +3458,26 @@ async fn remove_session_releases_workspace_binding_and_side_maps() { .expect("bind_local_session must succeed"); assert!(toolset_weak.upgrade().is_some()); *agent.workspace_ops.borrow_mut() = Some(ops); - agent.model_unavailable_sessions.borrow_mut().insert( - sid.0.to_string(), - acp::ModelId::new(std::sync::Arc::from("gone-model")), - ); + agent + .session_registry + .set_unavailable_model(&sid, acp::ModelId::new(std::sync::Arc::from("gone-model"))); agent.set_turn_number(&sid, 3); let (_permission_tx, permission_rx) = tokio::sync::mpsc::unbounded_channel::(); agent - .retained_resources - .borrow_mut() - .entry(sid.clone()) - .or_default() - .permission_event_receiver = Some(permission_rx); - agent - .resident_resources - .borrow_mut() - .entry(sid.clone()) - .or_default() - .require_gateway = true; + .session_registry + .set_permission_receiver(&sid, permission_rx); + agent.session_registry.mark_require_gateway(&sid); agent.remove_session(&sid); assert!( toolset_weak.upgrade().is_none(), "the workspace binding must release the toolset" ); - assert!( - !agent - .model_unavailable_sessions - .borrow() - .contains_key(sid.0.as_ref()) - ); - assert!(!agent.resident_resources.borrow().contains_key(&sid)); - assert!( - !agent.retained_resources.borrow().contains_key(&sid), + assert!(agent.session_registry.unavailable_model(&sid).is_none()); + assert_eq!(agent.session_registry.counts().resident_resources, 0); + assert_eq!( + agent.session_registry.counts().retained_resources, + 0, "retained per-session resources must be reclaimed on removal" ); } @@ -3484,6 +3755,46 @@ fn disconnect_keeps_resident_on_poisoned_lock() { ); }); } +/// A wedged actor stays tracked. `remove_session` releases everything else but +/// keeps a still-running thread, because dropping its handle would detach the +/// thread and leave nothing for the supervisor sweep to find. +#[test] +fn remove_session_keeps_a_running_thread_tracked() { + run_local_for_bridge_test(|| async { + let agent = build_minimal_agent_for_tests(); + let sid = acp::SessionId::new("sess-wedged"); + let (release_tx, release_rx) = std::sync::mpsc::channel::<()>(); + agent.session_registry.set_thread( + &sid, + crate::session::SessionThread::from_handle(std::thread::spawn(move || { + let _ = release_rx.recv(); + })), + ); + agent.set_turn_number(&sid, 1); + agent.remove_session(&sid); + assert!( + agent.session_registry.has_thread(&sid), + "a running actor thread must survive removal for the sweep" + ); + assert_eq!( + agent.session_registry.counts().retained_resources, + 0, + "everything except the running thread must be released" + ); + drop(release_tx); + for _ in 0..100 { + agent.sweep_dead_sessions(); + if !agent.session_registry.has_thread(&sid) { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + assert!( + !agent.session_registry.has_thread(&sid), + "the sweep must reclaim the thread once it exits" + ); + }); +} /// Idle-unload stub (memory bound) + supervisor interaction: a *fully idle* /// session is unloaded to disk on disconnect (actor `Shutdown`, handle /// dropped) while the `SessionThread` is **retained** for @@ -3498,8 +3809,8 @@ fn disconnect_unloads_idle_session_without_finalize() { agent.sessions.borrow_mut().insert(sid.clone(), handle); let mut observed = spawn_fake_actor(cmd_rx, false); let (release_tx, release_rx) = std::sync::mpsc::channel::<()>(); - agent.session_threads.borrow_mut().insert( - sid.clone(), + agent.session_registry.set_thread( + &sid, crate::session::SessionThread::from_handle(std::thread::spawn(move || { let _ = release_rx.recv(); })), @@ -3511,7 +3822,7 @@ fn disconnect_unloads_idle_session_without_finalize() { "idle session must be unloaded from the resident map on disconnect" ); assert!( - agent.session_threads.borrow().contains_key(&sid), + agent.session_registry.has_thread(&sid), "idle-unload must keep the SessionThread for reconnect drain" ); let shutdown = tokio::time::timeout(std::time::Duration::from_secs(1), observed.recv()) @@ -3534,13 +3845,13 @@ fn disconnect_unloads_idle_session_without_finalize() { drop(release_tx); let deadline = tokio::time::Instant::now() + (SESSION_SUPERVISOR_TICK * 6); while tokio::time::Instant::now() < deadline { - if !agent.session_threads.borrow().contains_key(&sid) { + if !agent.session_registry.has_thread(&sid) { break; } tokio::time::sleep(std::time::Duration::from_millis(20)).await; } assert!( - !agent.session_threads.borrow().contains_key(&sid), + !agent.session_registry.has_thread(&sid), "supervisor must drop the finished kept thread" ); assert!( @@ -3710,7 +4021,7 @@ fn session_live_state_map_is_bounded_across_cycles() { agent.close_session_explicit(&sid); } assert_eq!( - agent.session_live_state.borrow().len(), + agent.session_registry.counts().session_live_state, 0, "terminal closes must leave no residual live-state entries (bounded map)" ); @@ -3782,21 +4093,21 @@ fn supervisor_reaps_panicked_resident_actor() { let (handle, _tx, _rx) = make_live_session_handle(&sid, Some("turn-1")); agent.sessions.borrow_mut().insert(sid.clone(), handle); let panic_thread = std::thread::spawn(|| panic!("injected actor panic")); - agent.session_threads.borrow_mut().insert( - sid.clone(), + agent.session_registry.set_thread( + &sid, crate::session::SessionThread::from_handle(panic_thread), ); agent.set_session_live_state(&sid, SessionLiveState::Working); agent.ensure_session_supervisor(); let deadline = tokio::time::Instant::now() + (SESSION_SUPERVISOR_TICK * 6); while tokio::time::Instant::now() < deadline { - if !agent.session_threads.borrow().contains_key(&sid) { + if !agent.session_registry.has_thread(&sid) { break; } tokio::time::sleep(std::time::Duration::from_millis(20)).await; } assert!( - !agent.session_threads.borrow().contains_key(&sid), + !agent.session_registry.has_thread(&sid), "supervisor must reap the dead thread" ); assert!( diff --git a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests/dhat_soak.rs b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests/dhat_soak.rs index 70d60b6..6ab8a4e 100644 --- a/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests/dhat_soak.rs +++ b/crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests/dhat_soak.rs @@ -37,17 +37,11 @@ fn populate_and_evict(agent: &MvpAgent, i: usize) { } let (_ptx, prx) = tokio::sync::mpsc::unbounded_channel::(); - agent - .retained_resources - .borrow_mut() - .entry(sid.clone()) - .or_default() - .permission_event_receiver = Some(prx); + agent.session_registry.set_permission_receiver(&sid, prx); agent.set_turn_number(&sid, i as u64); - agent.model_unavailable_sessions.borrow_mut().insert( - sid.0.to_string(), - acp::ModelId::new(std::sync::Arc::from("gone-model")), - ); + agent + .session_registry + .set_unavailable_model(&sid, acp::ModelId::new(std::sync::Arc::from("gone-model"))); agent.remove_session(&sid); } diff --git a/crates/codegen/xai-grok-shell/src/auth/device_code.rs b/crates/codegen/xai-grok-shell/src/auth/device_code.rs index 8d21e3a..8c2bd4e 100644 --- a/crates/codegen/xai-grok-shell/src/auth/device_code.rs +++ b/crates/codegen/xai-grok-shell/src/auth/device_code.rs @@ -21,6 +21,11 @@ const DEFAULT_DEVICE_POLL_INTERVAL_SECS: i32 = 5; const DEVICE_SLOW_DOWN_INCREMENT_SECS: u64 = 5; const MIN_DEVICE_CODE_EXPIRY_FALLBACK_SECS: i64 = 10 * 60; +/// Only the 404 "no device endpoint" case is typed, because the login flow +/// matches on it to fall back to loopback. Every other device-code failure +/// stays a plain `anyhow` error: wrapping one in a `#[error(transparent)]` +/// variant hides the `reqwest::Error` the login funnel classifies, because +/// transparent forwards `source()` past the error it wraps. #[derive(Debug, Error)] pub enum DeviceCodeError { #[error( @@ -28,14 +33,6 @@ pub enum DeviceCodeError { Try `grok login` or set XAI_API_KEY instead." )] NotEnabled, - #[error(transparent)] - Other(#[from] anyhow::Error), -} - -impl From for DeviceCodeError { - fn from(e: reqwest::Error) -> Self { - Self::Other(e.into()) - } } // --- Public types --- @@ -137,7 +134,7 @@ pub async fn request_device_code( client_id: &str, scopes: &[String], surface: ClientSurface, -) -> Result { +) -> anyhow::Result { let client = crate::http::shared_client(); let url = format!("{}/oauth2/device/code", issuer.trim_end_matches('/')); let scope_str = scopes.join(" "); @@ -164,9 +161,9 @@ pub async fn request_device_code( let status = resp.status(); let body = resp.text().await.unwrap_or_default(); if status.as_u16() == 404 { - return Err(DeviceCodeError::NotEnabled); + anyhow::bail!(DeviceCodeError::NotEnabled); } - return Err(anyhow::anyhow!("Device code request failed (HTTP {status}): {body}").into()); + anyhow::bail!("Device code request failed (HTTP {status}): {body}"); } let server_resp: DeviceCodeResponse = resp.json().await?; @@ -177,10 +174,7 @@ pub async fn request_device_code( .chars() .all(|c| c.is_ascii_alphanumeric() || c == '-') { - return Err(anyhow::anyhow!( - "Server returned invalid user_code format (expected [A-Z0-9-])" - ) - .into()); + anyhow::bail!("Server returned invalid user_code format (expected [A-Z0-9-])"); } validate_verification_uri(&server_resp.verification_uri)?; diff --git a/crates/codegen/xai-grok-shell/src/auth/flow.rs b/crates/codegen/xai-grok-shell/src/auth/flow.rs index 4a33050..91f5cc7 100644 --- a/crates/codegen/xai-grok-shell/src/auth/flow.rs +++ b/crates/codegen/xai-grok-shell/src/auth/flow.rs @@ -7,7 +7,9 @@ use tokio::sync::{mpsc, oneshot}; use crate::auth::config::LEGACY_AUTH_SCOPE; use crate::auth::{AuthManager, GrokAuth, GrokComConfig, parse_output}; +use crate::http::TransportFailureKind; use crate::util::grok_home; +use xai_grok_telemetry::events::{LoginFailed, LoginFailureKind}; pub type StderrCallback = Box; @@ -451,6 +453,9 @@ pub async fn run_auth_flow_interactive( .await } +/// Every interactive login returns through here, so reporting the failure here +/// costs one event per attempt — a retried request, or the discovery cache +/// background token refresh shares, can't inflate it. Never changes the result. async fn run_auth_flow_inner( auth_manager: &Arc, grok_com_config: &GrokComConfig, @@ -460,6 +465,64 @@ async fn run_auth_flow_inner( url_tx: Option>>>>, code_rx: Option>, login_override: LoginTransportOverride, +) -> anyhow::Result<(GrokAuth, bool)> { + let result = run_auth_flow_steps( + auth_manager, + grok_com_config, + reauth, + force_interactive, + on_stderr, + url_tx, + code_rx, + login_override, + ) + .await; + if let Err(err) = &result + && let Some(event) = login_failure_event(err) + { + xai_grok_telemetry::session_ctx::log_event(event); + } + result +} + +/// `None` when nothing in the chain failed over HTTP (the user backed out, the +/// loopback listener couldn't bind, the id_token didn't validate) rather than +/// inventing a transport verdict for it. +fn login_failure_event(err: &anyhow::Error) -> Option { + let source = err + .chain() + .find_map(|cause| cause.downcast_ref::())?; + Some(LoginFailed { + error_kind: failure_kind( + crate::http::TransportFailure::classify(source).kind, + source.is_decode(), + ), + os_error: crate::http::find_os_error_code(source), + }) +} + +/// A body that won't parse is a decode failure, not a transport one — even +/// though `reqwest` also reports it as a body-phase error. +fn failure_kind(transport: TransportFailureKind, is_decode: bool) -> LoginFailureKind { + if is_decode { + return LoginFailureKind::Decode; + } + match transport { + TransportFailureKind::Unreachable => LoginFailureKind::TransportConnect, + TransportFailureKind::Interrupted => LoginFailureKind::TransportInterrupted, + TransportFailureKind::Permanent => LoginFailureKind::TransportPermanent, + } +} + +async fn run_auth_flow_steps( + auth_manager: &Arc, + grok_com_config: &GrokComConfig, + reauth: bool, + force_interactive: bool, + on_stderr: Option, + url_tx: Option>>>>, + code_rx: Option>, + login_override: LoginTransportOverride, ) -> anyhow::Result<(GrokAuth, bool)> { tracing::info!( has_oidc = grok_com_config.oidc.is_some(), @@ -900,6 +963,41 @@ pub async fn run_cli_login( oauth: bool, device_auth: bool, devbox: bool, +) -> anyhow::Result<()> { + // Devbox never reaches the login funnel, so it reports nothing and needs + // no telemetry client — and `AuthManager::new` is not free (it logs, and + // may rewrite auth.json to drop a stale scope). + if devbox { + let auth = super::devbox_login::run_devbox_login(config).await?; + return apply_post_login_config(auth).await; + } + + // Agent bootstrap is what normally initializes the product telemetry + // client, and `grok login` never boots an agent, so without this every + // event this process emits is dropped before reaching a sink. One manager + // serves both the identity it reads and the login flow below. + let auth_manager = Arc::new(AuthManager::new( + &grok_home::grok_home(), + config.grok_com_config.clone(), + )); + crate::agent::init::update_telemetry_config(config, &auth_manager); + + let result = run_cli_login_steps(config, &auth_manager, oauth, device_auth).await; + + // Posts run on a spawned task and this process exits as soon as we return. + xai_grok_telemetry::session_ctx::drain_pending(CLI_TELEMETRY_DRAIN).await; + result +} + +/// Returns as soon as the post lands (~1.7s cold), so the bound only bites on a +/// black-holed network — where waiting out the HTTP client timeout would be worse. +const CLI_TELEMETRY_DRAIN: std::time::Duration = std::time::Duration::from_secs(5); + +async fn run_cli_login_steps( + config: &crate::agent::config::Config, + auth_manager: &Arc, + oauth: bool, + device_auth: bool, ) -> anyhow::Result<()> { let login_override = LoginTransportOverride::from_flags(oauth, device_auth); @@ -908,15 +1006,11 @@ pub async fn run_cli_login( // supports the device flow. Without this guard, `grok login` on an // enterprise-OIDC deployment would wrongly enter the device branch (which // requires `oauth2`) and error. - let authenticated = if devbox { - super::devbox_login::run_devbox_login(config).await? - } else if cli_should_use_device(&config.grok_com_config, login_override).await { + let authenticated = if cli_should_use_device(&config.grok_com_config, login_override).await { if config.grok_com_config.oauth2.is_none() { // No OIDC and no oauth2 here, so `--oauth` can't help. anyhow::bail!("Sign-in is not available for this deployment. Set XAI_API_KEY instead."); } - let grok_home = grok_home::grok_home(); - let auth_manager = Arc::new(AuthManager::new(&grok_home, config.grok_com_config.clone())); // Route through the shared inner flow (not `run_device_code_login` // directly) so the external auth provider and devbox auto-migration run // before the interactive device login. `force_interactive` skips the @@ -925,7 +1019,7 @@ pub async fn run_cli_login( // Already resolved/logged above; pass `Preresolved(true)` so the inner flow // honors device without a second fetch or a duplicate `cli`-attributed log. let (auth, did_auth) = run_auth_flow_interactive( - &auth_manager, + auth_manager, &config.grok_com_config, None, None, @@ -945,21 +1039,35 @@ pub async fn run_cli_login( ); } // Loopback. `reauth=true` clears creds up front (legacy-scope hygiene), - // so abandoning logs you out — unlike the device branch above. + // so abandoning logs you out — unlike the device branch above. Calls + // `run_auth_flow` rather than `ensure_authenticated_with_override`, + // which would build a second `AuthManager`; with `reauth` set and no + // message prefix, the rest of that wrapper is a no-op. // Already resolved/logged above; pass `Preresolved(false)` so the inner // flow honors loopback without a duplicate `cli`-attributed log. - ensure_authenticated_with_override( + let (auth, did_auth) = run_auth_flow( + auth_manager, &config.grok_com_config, true, None, + None, + None, LoginTransportOverride::Preresolved(false), ) - .await? + .await?; + if did_auth { + report_signed_in(&auth); + } + auth }; - // Sync this principal's config now rather than waiting for the background - // tick. Stay quiet about absence/failure during login — confirm only when - // config was actually applied; `grok setup` reports the no-config case. + apply_post_login_config(authenticated).await +} + +/// Sync this principal's config now rather than waiting for the background +/// tick. Stay quiet about absence/failure during login — confirm only when +/// config was actually applied; `grok setup` reports the no-config case. +async fn apply_post_login_config(authenticated: GrokAuth) -> anyhow::Result<()> { let outcome = crate::managed_config::post_login_sync(Some(authenticated)).await; match outcome { crate::managed_config::ManagedConfigSync::Updated { is_team: true } => { @@ -1070,6 +1178,42 @@ mod tests { use crate::env::EnvVarGuard; use chrono::Utc; + /// `os_error` and the reqwest classification are covered in + /// `xai-grok-http`; what's local is which `LoginFailureKind` each maps to, + /// and that a decode failure never reads as a transport one. + #[test] + fn failure_kinds_map_one_to_one() { + assert_eq!( + failure_kind(TransportFailureKind::Unreachable, false), + LoginFailureKind::TransportConnect + ); + assert_eq!( + failure_kind(TransportFailureKind::Interrupted, false), + LoginFailureKind::TransportInterrupted + ); + assert_eq!( + failure_kind(TransportFailureKind::Permanent, false), + LoginFailureKind::TransportPermanent + ); + assert_eq!( + failure_kind(TransportFailureKind::Interrupted, true), + LoginFailureKind::Decode + ); + } + + /// A login that never reached the network is not a transport failure. The + /// positive path needs a real `reqwest::Error`, and building a client here + /// flips `jsonwebtoken` into its "no CryptoProvider" panic and breaks + /// unrelated auth tests, so classification is covered in `xai-grok-http`. + #[test] + fn non_http_login_failures_are_not_reported() { + let abandoned = anyhow::anyhow!("Login timed out after 10 minutes. Please try again."); + assert!(login_failure_event(&abandoned).is_none()); + + let nested = abandoned.context("Login failed. Please try again."); + assert!(login_failure_event(&nested).is_none()); + } + /// Run `f` with `GROK_LOGIN_DEVICE_FLOW` set to `value` (unset for `None`). /// `EnvVarGuard` serializes the process env and restores it on drop, so /// `resolve_device_flow` reads the env tier from a known state. diff --git a/crates/codegen/xai-grok-shell/src/auth/manager.rs b/crates/codegen/xai-grok-shell/src/auth/manager.rs index 9727269..355143d 100644 --- a/crates/codegen/xai-grok-shell/src/auth/manager.rs +++ b/crates/codegen/xai-grok-shell/src/auth/manager.rs @@ -18,7 +18,9 @@ pub(super) mod lock; mod sleep_gate; use lock::try_lock_auth_file_async; -use sleep_gate::{GateRaise, InFlightGuard, SleepGate}; +use sleep_gate::{InFlightGuard, SleepGate}; + +use crate::util::dual_clock::DualClock; use crate::auth::config::GrokComConfig; use crate::auth::error::AuthError; @@ -98,13 +100,13 @@ const RELOAD_RETRY_BACKOFF: StdDuration = StdDuration::from_millis(50); struct ScopedRefreshFailure { token_key: String, error: crate::auth::error::RefreshTokenFailedError, - /// Two-clock timestamp (see [`GateRaise`]): the TTL below is *real* time, + /// Two-clock timestamp (see [`DualClock`]): the TTL below is *real* time, /// so it must keep counting across a system sleep. The monotonic clock /// pauses during suspend — with it alone, a failure cached just before /// sleep would still short-circuit `auth()` for a further /// [`PERMANENT_FAILURE_TTL`] of *awake* time after wake, exactly when the /// user comes back and expects a recovered session. - recorded_at: GateRaise, + recorded_at: DualClock, } /// Auto-expiry safety net for the recoverable reasons (`ClientRejected`, @@ -202,11 +204,11 @@ pub struct AuthManager { /// manager so repeated 401s on the most-recent dead credential emit once. manual_auth: crate::auth::recovery::ManualAuthTracker, /// When the current unbroken run of dark-wake refresh deferrals began, on - /// two clocks (see [`GateRaise`]); `None` outside such a run. Bounds the + /// two clocks (see [`DualClock`]); `None` outside such a run. Bounds the /// deferral to [`sleep_gate::DARK_WAKE_DEFER_MAX`] so a machine stuck /// reporting dark wake can't defer refresh forever — see /// [`AuthManager::should_defer_for_dark_wake`]. - dark_wake_defer_since: parking_lot::RwLock>, + dark_wake_defer_since: parking_lot::RwLock>, /// Test-only override for [`AuthManager::is_dark_wake`]. `Some(_)` forces /// the dark-wake decision so the refresh-deferral path is unit-testable /// without a real macOS dark wake. `None` = consult the OS. @@ -2085,7 +2087,7 @@ impl AuthManager { *self.permanent_failure.write() = Some(ScopedRefreshFailure { token_key, error, - recorded_at: GateRaise::now(), + recorded_at: DualClock::now(), }); } @@ -2117,7 +2119,7 @@ impl AuthManager { /// on disk) must be allowed to refresh — otherwise a hard-expired sibling /// AT strands a process that could still refresh a live RT. /// - /// TTL expiry is judged on *both* clocks (see [`GateRaise`]): the monotonic + /// TTL expiry is judged on *both* clocks (see [`DualClock`]): the monotonic /// clock pauses during a system suspend, so a wall-clock arm is required /// for the TTL to elapse across sleep. Without it, a recoverable failure /// cached just before the lid closes (e.g. a transient escalation while @@ -2206,7 +2208,7 @@ impl AuthManager { // which the asserting test will surface loudly. let now_mono = std::time::Instant::now(); let now_wall = std::time::SystemTime::now(); - pf.recorded_at = GateRaise { + pf.recorded_at = DualClock { mono: now_mono.checked_sub(past_ttl).unwrap_or(now_mono), wall: now_wall.checked_sub(past_ttl).unwrap_or(now_wall), }; diff --git a/crates/codegen/xai-grok-shell/src/auth/manager/sleep_gate.rs b/crates/codegen/xai-grok-shell/src/auth/manager/sleep_gate.rs index da3c975..3835e1b 100644 --- a/crates/codegen/xai-grok-shell/src/auth/manager/sleep_gate.rs +++ b/crates/codegen/xai-grok-shell/src/auth/manager/sleep_gate.rs @@ -21,11 +21,12 @@ use std::sync::Arc; use std::sync::atomic::Ordering; -use std::time::{Duration as StdDuration, Instant, SystemTime}; +use std::time::{Duration as StdDuration, Instant}; use parking_lot::RwLock; use super::AuthManager; +use crate::util::dual_clock::DualClock; /// Max lifetime of the "system sleep imminent" gate. A wake event normally /// clears it; this is the safety bound so a *missed* wake event can never @@ -40,7 +41,7 @@ pub(super) const SLEEP_GATE_MAX: StdDuration = StdDuration::from_secs(120); /// interactive Mac with no display, whose system video capability is never set /// — which would otherwise defer every refresh forever and reach the same /// logged-out state this guard prevents. Bounded on two clocks (see -/// [`GateRaise`]) so it also survives the machine sleeping between dark wakes. +/// [`DualClock`]) so it also survives the machine sleeping between dark wakes. /// /// The straddle risk of one forced refresh is far smaller than a guaranteed /// logout: requests only force through while the machine is busy enough to @@ -65,57 +66,23 @@ pub(super) const SLEEP_ACK_MAX_WAIT: StdDuration = StdDuration::from_secs(20); #[cfg(not(target_os = "macos"))] pub(super) const SLEEP_ACK_MAX_WAIT: StdDuration = StdDuration::from_secs(3); -/// When a gate was raised, captured on *two* clocks so the [`SLEEP_GATE_MAX`] -/// backstop survives a system sleep. -/// -/// `Instant` is monotonic but, on macOS (`mach_absolute_time`) and Linux -/// (`CLOCK_MONOTONIC`), *pauses while the machine is asleep*. A gate raised just -/// before a long sleep would therefore never auto-expire on the monotonic clock -/// alone — the exact bug that let an expired token reach the server and 401. -/// The wall clock (`SystemTime`) keeps advancing through sleep, so we expire the -/// gate once *either* clock passes the bound: -/// - the monotonic clock bounds elapsed *awake* time (immune to wall-clock -/// jumps from NTP / manual changes), and -/// - the wall clock bounds elapsed *real* time (immune to the sleep pause). -#[derive(Clone, Copy)] -pub(super) struct GateRaise { - /// Monotonic; pauses during sleep. Bounds elapsed *awake* time. - pub(super) mono: Instant, - /// Wall clock; advances through sleep. Bounds elapsed *real* time. - pub(super) wall: SystemTime, -} - -impl GateRaise { - pub(super) fn now() -> Self { - Self { - mono: Instant::now(), - wall: SystemTime::now(), - } - } - - /// Elapsed on each clock as `(monotonic, wall)`. Wall-clock elapsed is - /// clamped to zero if the clock ran backwards (NTP step / manual change) so - /// a backward jump can never *extend* the gate — the monotonic clock still - /// bounds it in that case. - pub(super) fn elapsed(&self) -> (StdDuration, StdDuration) { - ( - self.mono.elapsed(), - self.wall.elapsed().unwrap_or(StdDuration::ZERO), - ) - } -} - /// A gate `refresh_chain` consults to avoid *starting* an IdP refresh just /// before sleep. Only *defers* a not-yet-started refresh; an in-flight one is /// left to finish (see [`AuthManager::refresh_chain`]). +/// +/// The raise timestamp is a [`DualClock`] so the [`SLEEP_GATE_MAX`] backstop +/// survives the sleep itself: a gate raised just before a long sleep would +/// never auto-expire on the monotonic clock alone — the exact bug that let +/// an expired token reach the server and 401 — so the gate expires once +/// *either* clock passes the bound. #[derive(Default)] pub(super) struct SleepGate { - pub(super) raised_at: RwLock>, + pub(super) raised_at: RwLock>, } impl SleepGate { pub(super) fn raise(&self) { - *self.raised_at.write() = Some(GateRaise::now()); + *self.raised_at.write() = Some(DualClock::now()); xai_grok_telemetry::unified_log::warn("auth.sleep.gate_set", None, None); } @@ -142,7 +109,7 @@ impl SleepGate { /// A stale gate (a missed/late wake event) is lazily lowered here so it can /// never permanently block refresh; this read can therefore have a side /// effect. The gate expires once *either* clock passes [`SLEEP_GATE_MAX`] - /// (see [`GateRaise`]): without the wall-clock arm, a gate raised before a + /// (see [`DualClock`]): without the wall-clock arm, a gate raised before a /// long sleep would never auto-expire, because the monotonic clock pauses /// while the machine is asleep. pub(super) fn is_gated(&self) -> bool { @@ -355,7 +322,7 @@ impl AuthManager { /// in a dark wake — bounded so deferral can never be indefinite. /// /// Tracks when the current unbroken run of dark-wake deferrals began (on two - /// clocks; see [`GateRaise`]). While inside the [`DARK_WAKE_DEFER_MAX`] + /// clocks; see [`DualClock`]). While inside the [`DARK_WAKE_DEFER_MAX`] /// budget it returns `true` (defer). Once either clock passes the bound it /// forces one refresh through (`false`) and resets the clock, so a machine /// stuck reporting a continuous dark wake refreshes periodically instead of @@ -375,7 +342,7 @@ impl AuthManager { } let Some(raise) = *run else { // First deferral of this dark-wake run: start the budget clock. - *run = Some(GateRaise::now()); + *run = Some(DualClock::now()); return true; }; let (mono, wall) = raise.elapsed(); diff --git a/crates/codegen/xai-grok-shell/src/auth/manager_tests.rs b/crates/codegen/xai-grok-shell/src/auth/manager_tests.rs index 7034df6..4dc22af 100644 --- a/crates/codegen/xai-grok-shell/src/auth/manager_tests.rs +++ b/crates/codegen/xai-grok-shell/src/auth/manager_tests.rs @@ -4353,7 +4353,7 @@ async fn dark_wake_defer_forces_refresh_after_max() { ) else { return; // machine/clock can't represent the backdate — skip }; - *mgr.dark_wake_defer_since.write() = Some(super::sleep_gate::GateRaise { mono, wall }); + *mgr.dark_wake_defer_since.write() = Some(crate::util::dual_clock::DualClock { mono, wall }); assert_eq!( mgr.auth().await.unwrap().key, @@ -4490,7 +4490,7 @@ async fn sleep_gate_auto_expires_after_max() { ) else { return; // machine/clock can't represent the backdate — not reproducible; skip }; - *mgr.sleep_gate.raised_at.write() = Some(super::sleep_gate::GateRaise { mono, wall }); + *mgr.sleep_gate.raised_at.write() = Some(crate::util::dual_clock::DualClock { mono, wall }); assert!( !mgr.is_sleep_gated(), @@ -4522,7 +4522,7 @@ async fn sleep_gate_auto_expires_when_wall_clock_passes_during_sleep() { let Some(wall) = std::time::SystemTime::now().checked_sub(back) else { return; // clock can't represent the backdate — not reproducible; skip }; - *mgr.sleep_gate.raised_at.write() = Some(super::sleep_gate::GateRaise { + *mgr.sleep_gate.raised_at.write() = Some(crate::util::dual_clock::DualClock { mono: Instant::now(), wall, }); diff --git a/crates/codegen/xai-grok-shell/src/config/watcher.rs b/crates/codegen/xai-grok-shell/src/config/watcher.rs index 4364396..dfe7987 100644 --- a/crates/codegen/xai-grok-shell/src/config/watcher.rs +++ b/crates/codegen/xai-grok-shell/src/config/watcher.rs @@ -399,12 +399,6 @@ fn log_watch_error(err: ¬ify::Error, msg: &str) { } } -pub struct SkillsFileWatcher { - debouncer: Debouncer, - refresh_dirs: Vec<(PathBuf, RecursiveMode)>, - refreshed_dirs: HashSet, -} - const SKILLS_DEBOUNCE: Duration = Duration::from_secs(2); #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -414,17 +408,18 @@ pub enum DiscoveryChange { } fn discovery_change_for_path(path: &Path) -> Option { - if path.file_name().is_some_and(|name| name == ".grok") { + let file_name = path.file_name().and_then(|name| name.to_str()); + if file_name.is_some_and(|name| VENDOR_CONFIG_ROOT_NAMES.contains(&name)) { return Some(DiscoveryChange::Skills); } - if path.file_name().is_some_and(|name| name == "workflows") + if file_name.is_some_and(|name| name == "workflows") || path .ancestors() .any(|ancestor| ancestor.file_name().is_some_and(|name| name == "workflows")) { return Some(DiscoveryChange::Workflows); } - if path.file_name().is_some_and(|name| name == "SKILL.md") + if file_name.is_some_and(|name| name == "skills" || name == "commands" || name == "SKILL.md") || path .ancestors() .any(|ancestor| ancestor.file_name().is_some_and(|name| name == "skills")) @@ -438,72 +433,158 @@ fn discovery_change_for_path(path: &Path) -> Option { None } -/// True for a global/home-level config dir that must never be watched -/// recursively: `grok_home` (`~/.grok`, or `$GROK_HOME`) or a known vendor dir -/// directly under `$HOME` ([`HOME_VENDOR_DIRS`]). -/// -/// These hold large non-skill trees — `~/.grok` alone has `worktrees/`, -/// `sessions/`, `logs/`, `upload_queue/` — so recursing them exhausted the -/// inotify quota (~780k watches on a devbox) and, since each worktree is a full -/// checkout, fired skill reloads on ordinary repo activity. They get scoped -/// watches instead ([`watch_skill_subdirs`]); project/repo dirs — and -/// user-supplied `[skills].paths` entries, which discovery walks in full — stay -/// recursive. Matching only these specific names (not "any dir whose parent is -/// `$HOME`") is what keeps a `[skills].paths = ["~/my-skills"]` fully watched. -fn is_global_config_dir(dir: &Path, grok_home: &Path) -> bool { - #[allow(deprecated)] - let home = std::env::home_dir(); - is_global_config_dir_impl(dir, grok_home, home.as_deref()) -} +/// Known vendor config root basenames; kept in sync with `collect_skill_config_dirs`. +const VENDOR_CONFIG_ROOT_NAMES: &[&str] = &[".grok", ".agents", ".claude", ".cursor"]; -/// Vendor config dir names that sit directly under `$HOME` and carry large -/// non-skill trees. Kept in sync with the home-level dirs added by -/// `collect_skill_config_dirs`. -const HOME_VENDOR_DIRS: &[&str] = &[".grok", ".agents", ".claude", ".cursor"]; - -/// Testable core of [`is_global_config_dir`] with `$HOME` injected. -fn is_global_config_dir_impl(dir: &Path, grok_home: &Path, home: Option<&Path>) -> bool { - let canon = |p: &Path| dunce::canonicalize(p).unwrap_or_else(|_| p.to_path_buf()); - if canon(dir) == canon(grok_home) { +/// Vendor roots (by name or `grok_home`) must use scoped watches — they can +/// contain large non-skill trees (`worktrees/`, etc.). +fn is_vendor_config_root(dir: &Path, grok_home: &Path) -> bool { + if paths_equal(dir, grok_home) { return true; } - let Some(home) = home else { return false }; - if dir.parent().map(canon) != Some(canon(home)) { - return false; - } dir.file_name() .and_then(|n| n.to_str()) - .is_some_and(|n| HOME_VENDOR_DIRS.contains(&n)) + .is_some_and(|n| VENDOR_CONFIG_ROOT_NAMES.contains(&n)) } +fn paths_equal(a: &Path, b: &Path) -> bool { + if a == b { + return true; + } + let canon = |p: &Path| dunce::canonicalize(p).unwrap_or_else(|_| p.to_path_buf()); + canon(a) == canon(b) +} + +fn dirs_contain(dirs: &[PathBuf], target: &Path) -> bool { + dirs.iter().any(|dir| paths_equal(dir, target)) +} + +fn path_set_contains(dirs: &HashSet, target: &Path) -> bool { + dirs.iter().any(|dir| paths_equal(dir, target)) +} + +fn vendor_skill_refresh_dirs(config_dir: &Path) -> [(PathBuf, RecursiveMode); 3] { + [ + (config_dir.join("skills"), RecursiveMode::Recursive), + (config_dir.join("commands"), RecursiveMode::NonRecursive), + (config_dir.join("workflows"), RecursiveMode::NonRecursive), + ] +} + +fn project_grok_refresh_dirs(project_root: &Path) -> Vec<(PathBuf, RecursiveMode)> { + let project_grok = project_root.join(".grok"); + let mut dirs = vec![(project_grok.clone(), RecursiveMode::NonRecursive)]; + dirs.extend(vendor_skill_refresh_dirs(&project_grok)); + dirs +} + +fn attach_new_refresh_dirs( + debouncer: &mut Debouncer, + refresh_dirs: &[(PathBuf, RecursiveMode)], + refreshed_dirs: &mut HashSet, + err_msg: &str, +) -> bool { + let mut changed = false; + for (dir, mode) in refresh_dirs { + if path_set_contains(refreshed_dirs, dir) || !dir.is_dir() { + continue; + } + match debouncer.watcher().watch(dir, *mode) { + Ok(()) => { + refreshed_dirs.insert(dir.clone()); + changed = true; + } + Err(error) => log_watch_error(&error, err_msg), + } + } + changed +} + +/// Paths successfully watched under a scoped vendor root (root + skill subdirs). fn watch_skill_subdirs( debouncer: &mut Debouncer, config_dir: &Path, -) -> usize { - let mut watched = 0; +) -> HashSet { + let mut watched = HashSet::new(); match debouncer .watcher() .watch(config_dir, RecursiveMode::NonRecursive) { - Ok(()) => watched += 1, + Ok(()) => { + watched.insert(config_dir.to_path_buf()); + } Err(error) => log_watch_error(&error, "failed to watch config dir root"), } - for (subdir, mode) in [ - ("skills", RecursiveMode::Recursive), - ("commands", RecursiveMode::NonRecursive), - ("workflows", RecursiveMode::NonRecursive), - ] { - let dir = config_dir.join(subdir); - if dir.is_dir() { - match debouncer.watcher().watch(&dir, mode) { - Ok(()) => watched += 1, - Err(error) => log_watch_error(&error, "failed to watch discovery subdir"), + for (dir, mode) in vendor_skill_refresh_dirs(config_dir) { + if !dir.is_dir() { + continue; + } + match debouncer.watcher().watch(&dir, mode) { + Ok(()) => { + watched.insert(dir); } + Err(error) => log_watch_error(&error, "failed to watch discovery subdir"), } } watched } +#[derive(Debug, Clone, PartialEq, Eq)] +struct SkillsWatchPlan { + vendor_roots: Vec, + recursive_roots: Vec, + /// Non-recursive parent so first create of a missing project vendor root is observed. + project_parent_watch: Option, + refresh_dirs: Vec<(PathBuf, RecursiveMode)>, +} + +/// Pure composition: classify discovery roots and seed mid-session refresh targets. +fn plan_skills_watch_targets( + dirs_to_watch: &[PathBuf], + grok_home: &Path, + project_root: Option<&Path>, +) -> SkillsWatchPlan { + let mut vendor_roots = Vec::new(); + let mut recursive_roots = Vec::new(); + let mut refresh_dirs = Vec::new(); + + for dir in dirs_to_watch { + if is_vendor_config_root(dir, grok_home) { + vendor_roots.push(dir.clone()); + refresh_dirs.extend(vendor_skill_refresh_dirs(dir)); + } else { + recursive_roots.push(dir.clone()); + } + } + + let mut project_parent_watch = None; + if let Some(project_root) = project_root { + let mut missing_project_vendor = false; + for name in VENDOR_CONFIG_ROOT_NAMES { + let vendor_root = project_root.join(name); + if !dirs_contain(dirs_to_watch, &vendor_root) { + missing_project_vendor = true; + refresh_dirs.push((vendor_root.clone(), RecursiveMode::NonRecursive)); + refresh_dirs.extend(vendor_skill_refresh_dirs(&vendor_root)); + } + } + if missing_project_vendor && !dirs_contain(dirs_to_watch, project_root) { + project_parent_watch = Some(project_root.to_path_buf()); + } + } + + SkillsWatchPlan { + vendor_roots, + recursive_roots, + project_parent_watch, + refresh_dirs, + } +} + +/// Watches project `.grok` skills/commands/workflows for mid-session discovery. +/// +/// After a [`DiscoveryChange`], call [`Self::refresh_new_dirs`] so newly created +/// seed dirs get watches attached. pub struct ProjectDiscoveryWatcher { debouncer: Debouncer, refresh_dirs: Vec<(PathBuf, RecursiveMode)>, @@ -514,7 +595,6 @@ impl ProjectDiscoveryWatcher { pub fn start(cwd: &Path) -> Option<(Self, mpsc::UnboundedReceiver)> { let project_root = crate::session::workflow::registry::project_root(cwd); let project_grok = project_root.join(".grok"); - let workflows = project_grok.join("workflows"); let (tx, rx) = mpsc::unbounded_channel(); let project_grok_for_events = project_grok.clone(); let mut debouncer = @@ -552,30 +632,14 @@ impl ProjectDiscoveryWatcher { log_watch_error(&error, "failed to watch project workflow parent"); return None; } - let refresh_dirs = vec![ - (project_grok, RecursiveMode::NonRecursive), - ( - project_root.join(".grok").join("skills"), - RecursiveMode::Recursive, - ), - ( - project_root.join(".grok").join("commands"), - RecursiveMode::NonRecursive, - ), - (workflows, RecursiveMode::NonRecursive), - ]; + let refresh_dirs = project_grok_refresh_dirs(&project_root); let mut refreshed_dirs = HashSet::from([initial]); - for (dir, mode) in &refresh_dirs { - if refreshed_dirs.contains(dir) || !dir.is_dir() { - continue; - } - match debouncer.watcher().watch(dir, *mode) { - Ok(()) => { - refreshed_dirs.insert(dir.clone()); - } - Err(error) => log_watch_error(&error, "failed to watch project discovery dir"), - } - } + attach_new_refresh_dirs( + &mut debouncer, + &refresh_dirs, + &mut refreshed_dirs, + "failed to watch project discovery dir", + ); Some(( Self { debouncer, @@ -586,32 +650,61 @@ impl ProjectDiscoveryWatcher { )) } + /// Attach watches for seed dirs that now exist (call after a discovery event). pub fn refresh_new_dirs(&mut self) { - for (dir, mode) in &self.refresh_dirs { - if self.refreshed_dirs.contains(dir) || !dir.is_dir() { - continue; - } - match self.debouncer.watcher().watch(dir, *mode) { - Ok(()) => { - self.refreshed_dirs.insert(dir.clone()); - } - Err(error) => { - log_watch_error(&error, "failed to watch newly-created project workflow dir") - } - } - } + attach_new_refresh_dirs( + &mut self.debouncer, + &self.refresh_dirs, + &mut self.refreshed_dirs, + "failed to watch newly-created project workflow dir", + ); } } +/// Watches skill/command/workflow discovery dirs and classifies disk changes. +pub struct SkillsFileWatcher { + debouncer: Debouncer, + refresh_dirs: Vec<(PathBuf, RecursiveMode)>, + refreshed_dirs: HashSet, +} + impl SkillsFileWatcher { + /// Start watching discovery dirs from + /// [`collect_skill_config_dirs`](xai_grok_agent::prompt::skills::collect_skill_config_dirs). /// - /// Uses [`collect_skill_config_dirs`](xai_grok_agent::prompt::skills::collect_skill_config_dirs) - /// as the canonical directory source so the watcher covers the same - /// locations as skill discovery. + /// After a [`DiscoveryChange`], call [`Self::refresh_new_discovery_dirs`] so + /// newly created seed dirs get watches attached. pub fn start( cwd: Option<&Path>, monorepo_user_dir: Option<&Path>, config_paths: &[String], + ) -> Option<(Self, mpsc::UnboundedReceiver)> { + let grok_home = xai_grok_tools::util::grok_home::grok_home(); + // Watch the full superset of vendor dirs (all-on compat). This watcher + // is leader-global (no per-session compat resolved here); the actual + // per-session discovery gating happens downstream, so watching a + // currently-disabled vendor dir is harmless (a change just re-runs the + // gated discovery) and avoids ever missing a watch if a toggle flips. + let dirs_to_watch = xai_grok_agent::prompt::skills::collect_skill_config_dirs( + cwd, + monorepo_user_dir, + &grok_home, + config_paths, + xai_grok_tools::types::compat::CompatConfig::default(), + ); + let project_root = cwd.map(crate::session::workflow::registry::project_root); + Self::start_with_dirs(&dirs_to_watch, &grok_home, project_root.as_deref()) + } + + /// Start with explicit discovery roots (benches and isolated tests). + /// + /// Production code should prefer [`Self::start`], which collects the same + /// dir set discovery uses. After a [`DiscoveryChange`], call + /// [`Self::refresh_new_discovery_dirs`]. + pub fn start_with_dirs( + dirs_to_watch: &[PathBuf], + grok_home: &Path, + project_root: Option<&Path>, ) -> Option<(Self, mpsc::UnboundedReceiver)> { let (tx, rx) = mpsc::unbounded_channel(); @@ -636,30 +729,37 @@ impl SkillsFileWatcher { .map_err(|e| tracing::warn!(error = %e, "failed to create skills file watcher")) .ok()?; - let grok_home = xai_grok_tools::util::grok_home::grok_home(); - // Watch the full superset of vendor dirs (all-on compat). This watcher - // is leader-global (no per-session compat resolved here); the actual - // per-session discovery gating happens downstream, so watching a - // currently-disabled vendor dir is harmless (a change just re-runs the - // gated discovery) and avoids ever missing a watch if a toggle flips. - let dirs_to_watch = xai_grok_agent::prompt::skills::collect_skill_config_dirs( - cwd, - monorepo_user_dir, - &grok_home, - config_paths, - xai_grok_tools::types::compat::CompatConfig::default(), - ); + let plan = plan_skills_watch_targets(dirs_to_watch, grok_home, project_root); + let mut watched = 0; - for dir in &dirs_to_watch { - if is_global_config_dir(dir, &grok_home) { - watched += watch_skill_subdirs(&mut debouncer, dir); - } else { - // Project/repo dir: bounded, so recurse to catch new `skills/` - // dirs created mid-session as well as edits to existing files. - match debouncer.watcher().watch(dir, RecursiveMode::Recursive) { - Ok(()) => watched += 1, - Err(e) => log_watch_error(&e, "failed to watch directory for skill changes"), + let mut refreshed_dirs = HashSet::new(); + for dir in &plan.vendor_roots { + let attached = watch_skill_subdirs(&mut debouncer, dir); + watched += attached.len(); + refreshed_dirs.extend(attached); + } + for dir in &plan.recursive_roots { + match debouncer.watcher().watch(dir, RecursiveMode::Recursive) { + Ok(()) => { + watched += 1; + refreshed_dirs.insert(dir.clone()); } + Err(e) => log_watch_error(&e, "failed to watch directory for skill changes"), + } + } + if let Some(parent_watch) = &plan.project_parent_watch { + match debouncer + .watcher() + .watch(parent_watch, RecursiveMode::NonRecursive) + { + Ok(()) => { + watched += 1; + refreshed_dirs.insert(parent_watch.clone()); + } + Err(error) => log_watch_error( + &error, + "failed to watch workflow discovery parent directory", + ), } } @@ -670,62 +770,25 @@ impl SkillsFileWatcher { tracing::info!(dirs = watched, "skills file watcher started"); - let mut refresh_dirs = vec![(grok_home.join("workflows"), RecursiveMode::NonRecursive)]; - if let Some(cwd) = cwd { - let project_root = crate::session::workflow::registry::project_root(cwd); - let project_grok = project_root.join(".grok"); - let parent_watch = if project_grok.is_dir() { - project_grok.clone() - } else { - project_root - }; - if !dirs_to_watch.iter().any(|dir| dir == &parent_watch) { - match debouncer - .watcher() - .watch(&parent_watch, RecursiveMode::NonRecursive) - { - Ok(()) => {} - Err(error) => log_watch_error( - &error, - "failed to watch workflow discovery parent directory", - ), - } - } - refresh_dirs.push((project_grok.clone(), RecursiveMode::NonRecursive)); - refresh_dirs.push((project_grok.join("workflows"), RecursiveMode::NonRecursive)); - } - let refreshed_dirs = refresh_dirs - .iter() - .filter(|(dir, _)| dir.is_dir()) - .map(|(dir, _)| dir.clone()) - .collect(); Some(( Self { debouncer, - refresh_dirs, + refresh_dirs: plan.refresh_dirs, refreshed_dirs, }, rx, )) } + /// Attach watches for seed dirs that now exist (call after a discovery event). + /// Returns true if any new watch was attached. pub fn refresh_new_discovery_dirs(&mut self) -> bool { - let mut changed = false; - for (dir, mode) in &self.refresh_dirs { - if self.refreshed_dirs.contains(dir) || !dir.is_dir() { - continue; - } - match self.debouncer.watcher().watch(dir, *mode) { - Ok(()) => { - self.refreshed_dirs.insert(dir.clone()); - changed = true; - } - Err(error) => { - log_watch_error(&error, "failed to watch newly-created discovery directory") - } - } - } - changed + attach_new_refresh_dirs( + &mut self.debouncer, + &self.refresh_dirs, + &mut self.refreshed_dirs, + "failed to watch newly-created discovery directory", + ) } } @@ -739,41 +802,389 @@ mod tests { std::thread::sleep(Duration::from_millis(ms)); } - /// `is_global_config_dir` must scope down only grok_home and the known - /// vendor dirs under `$HOME` — NOT arbitrary `[skills].paths` entries such - /// as `~/my-skills`, whose skills discovery walks in full and so must stay - /// recursively watched. #[test] - fn is_global_config_dir_matches_only_grok_home_and_vendor_dirs() { + fn is_vendor_config_root_matches_known_names_at_any_tier() { let home = TempDir::new().unwrap(); let home = home.path(); let grok_home = home.join(".grok"); - let g = |dir: &Path| is_global_config_dir_impl(dir, &grok_home, Some(home)); + assert!(is_vendor_config_root(&grok_home, &grok_home)); + assert!(is_vendor_config_root(&home.join(".claude"), &grok_home)); + assert!(is_vendor_config_root(&home.join(".cursor"), &grok_home)); + assert!(is_vendor_config_root(&home.join(".agents"), &grok_home)); + assert!(is_vendor_config_root( + &home.join("repo").join(".grok"), + &grok_home + )); + assert!(is_vendor_config_root( + &home.join("repo").join(".claude"), + &grok_home + )); - // grok_home and vendor dirs directly under $HOME: scoped (global). - assert!(g(&grok_home)); - assert!(g(&home.join(".claude"))); - assert!(g(&home.join(".cursor"))); - assert!(g(&home.join(".agents"))); + assert!(!is_vendor_config_root(&home.join("my-skills"), &grok_home)); + assert!(!is_vendor_config_root(&home.join(".config"), &grok_home)); + assert!(!is_vendor_config_root( + &home.join("repo").join("my-skills"), + &grok_home + )); - // A user [skills].paths entry under $HOME: NOT global (stays recursive). - assert!(!g(&home.join("my-skills"))); - assert!(!g(&home.join(".config"))); - // A project/repo config dir (parent isn't $HOME): NOT global. - assert!(!g(&home.join("repo").join(".grok"))); + let custom_home = home.join("custom-grok-home"); + assert!(is_vendor_config_root(&custom_home, &custom_home)); + } + + #[test] + fn vendor_skill_refresh_dirs_paths_and_modes() { + let root = Path::new("/tmp/project/.claude"); + assert_eq!( + vendor_skill_refresh_dirs(root), + [ + (root.join("skills"), RecursiveMode::Recursive), + (root.join("commands"), RecursiveMode::NonRecursive), + (root.join("workflows"), RecursiveMode::NonRecursive), + ] + ); + } + + #[test] + fn project_grok_refresh_dirs_matches_vendor_layout() { + let project = Path::new("/tmp/repo"); + let grok = project.join(".grok"); + let dirs = project_grok_refresh_dirs(project); + + assert_eq!(dirs.len(), 4); + assert_eq!(dirs[0], (grok.clone(), RecursiveMode::NonRecursive)); + assert_eq!( + &dirs[1..], + [ + (grok.join("skills"), RecursiveMode::Recursive), + (grok.join("commands"), RecursiveMode::NonRecursive), + (grok.join("workflows"), RecursiveMode::NonRecursive), + ] + ); + assert_eq!(dirs[1..], vendor_skill_refresh_dirs(&grok)); + } + + #[test] + #[cfg(unix)] + fn path_set_contains_uses_paths_equal() { + let tmp = TempDir::new().unwrap(); + let real = tmp.path().join("real_skills"); + fs::create_dir_all(&real).unwrap(); + let link = tmp.path().join("link_skills"); + std::os::unix::fs::symlink(&real, &link).unwrap(); + + assert_ne!(real.as_os_str(), link.as_os_str()); + let mut set = HashSet::new(); + set.insert(real.clone()); + + assert!(path_set_contains(&set, &real)); + assert!( + path_set_contains(&set, &link), + "symlink form must match via paths_equal/canonicalize" + ); + assert!( + !set.contains(&link), + "HashSet::contains must not match symlink form" + ); + assert!(!path_set_contains(&set, &tmp.path().join("other"))); + } + + #[test] + #[cfg(unix)] + fn attach_new_refresh_dirs_skips_known_and_missing() { + let tmp = TempDir::new().unwrap(); + let root = tmp.path(); + let known_real = root.join("known"); + let known_link = root.join("known_link"); + let missing = root.join("missing"); + let fresh = root.join("fresh"); + fs::create_dir(&known_real).unwrap(); + std::os::unix::fs::symlink(&known_real, &known_link).unwrap(); + fs::create_dir(&fresh).unwrap(); + + let mut debouncer = new_filtered_debouncer(Duration::from_millis(50), |_| {}).unwrap(); + debouncer + .watcher() + .watch(root, RecursiveMode::NonRecursive) + .unwrap(); + + // Seed with symlink form; refresh_dirs lists the real path (paths_equal, not byte-equal). + let refresh_dirs = vec![ + (known_real.clone(), RecursiveMode::NonRecursive), + (missing.clone(), RecursiveMode::NonRecursive), + (fresh.clone(), RecursiveMode::NonRecursive), + ]; + let mut refreshed_dirs = HashSet::from([known_link.clone()]); + assert_ne!(known_real.as_os_str(), known_link.as_os_str()); + assert!(!refreshed_dirs.contains(&known_real)); + + assert!(attach_new_refresh_dirs( + &mut debouncer, + &refresh_dirs, + &mut refreshed_dirs, + "test attach", + )); + assert_eq!( + refreshed_dirs.len(), + 2, + "skip known (path-equal form) and missing; only fresh attaches" + ); + assert!(path_set_contains(&refreshed_dirs, &known_real)); + assert!(path_set_contains(&refreshed_dirs, &known_link)); + assert!(path_set_contains(&refreshed_dirs, &fresh)); + assert!(!path_set_contains(&refreshed_dirs, &missing)); + assert!(!attach_new_refresh_dirs( + &mut debouncer, + &refresh_dirs, + &mut refreshed_dirs, + "test attach", + )); + assert_eq!(refreshed_dirs.len(), 2); + } + + fn expected_missing_vendor_refresh_seeds(project_root: &Path) -> Vec<(PathBuf, RecursiveMode)> { + let mut expected = Vec::new(); + for name in VENDOR_CONFIG_ROOT_NAMES { + let root = project_root.join(name); + expected.push((root.clone(), RecursiveMode::NonRecursive)); + expected.extend(vendor_skill_refresh_dirs(&root)); + } + expected + } + + /// Parent-only plan (empty dirs_to_watch) must still start so mid-session + /// vendor creates under project_root can attach via refresh seeds. + #[test] + fn start_with_dirs_keeps_parent_only_watch() { + let tmp = TempDir::new().unwrap(); + let project = tmp.path(); + let grok_home = project.join("home-grok"); + fs::create_dir_all(&grok_home).unwrap(); + + let plan = plan_skills_watch_targets(&[], &grok_home, Some(project)); + assert!(plan.vendor_roots.is_empty()); + assert!(plan.recursive_roots.is_empty()); + assert_eq!(plan.project_parent_watch.as_deref(), Some(project)); + + let (watcher, _rx) = SkillsFileWatcher::start_with_dirs(&[], &grok_home, Some(project)) + .expect("parent-only watch must start when no discovery roots exist yet"); + assert!( + path_set_contains(&watcher.refreshed_dirs, project), + "successful project parent watch must be retained" + ); + assert_eq!( + watcher.refresh_dirs, + expected_missing_vendor_refresh_seeds(project) + ); + } + + #[test] + fn plan_skills_watch_targets_scopes_vendors_and_seeds_refresh() { + let tmp = TempDir::new().unwrap(); + let project = tmp.path(); + let grok_home = project.join("home-grok"); + let project_claude = project.join(".claude"); + let project_grok = project.join(".grok"); + let custom = project.join("my-skills"); + fs::create_dir_all(&project_claude).unwrap(); + fs::create_dir_all(&project_grok).unwrap(); + fs::create_dir_all(&custom).unwrap(); + + let dirs = vec![project_claude.clone(), project_grok.clone(), custom.clone()]; + let plan = plan_skills_watch_targets(&dirs, &grok_home, Some(project)); + + assert_eq!( + plan.vendor_roots, + vec![project_claude.clone(), project_grok.clone()] + ); + assert_eq!(plan.recursive_roots, vec![custom]); + assert_eq!(plan.project_parent_watch.as_deref(), Some(project)); + + let mut expected_refresh: Vec<(PathBuf, RecursiveMode)> = + vendor_skill_refresh_dirs(&project_claude) + .into_iter() + .chain(vendor_skill_refresh_dirs(&project_grok)) + .collect(); + for name in [".agents", ".cursor"] { + let root = project.join(name); + expected_refresh.push((root.clone(), RecursiveMode::NonRecursive)); + expected_refresh.extend(vendor_skill_refresh_dirs(&root)); + } + assert_eq!(plan.refresh_dirs, expected_refresh); + } + + #[test] + fn plan_skills_watch_targets_multi_vendor_refresh_fanout() { + let grok_home = PathBuf::from("/home/u/.grok"); + let a = PathBuf::from("/repo/.claude"); + let b = PathBuf::from("/repo/.agents"); + let plan = plan_skills_watch_targets(&[a.clone(), b.clone()], &grok_home, None); + + assert_eq!(plan.vendor_roots, vec![a.clone(), b.clone()]); + assert!(plan.recursive_roots.is_empty()); + assert_eq!( + plan.refresh_dirs, + vendor_skill_refresh_dirs(&a) + .into_iter() + .chain(vendor_skill_refresh_dirs(&b)) + .collect::>() + ); + for root in [&a, &b] { + for (dir, mode) in vendor_skill_refresh_dirs(root) { + assert!( + plan.refresh_dirs.contains(&(dir, mode)), + "missing refresh seed for {}", + root.display() + ); + } + } + } + + #[test] + fn plan_skills_watch_targets_seeds_all_missing_project_vendor_roots() { + let tmp = TempDir::new().unwrap(); + let project = tmp.path(); + let grok_home = project.join("elsewhere").join(".grok"); + let plan = plan_skills_watch_targets(&[], &grok_home, Some(project)); + + assert_eq!(plan.project_parent_watch.as_deref(), Some(project)); + assert!(plan.vendor_roots.is_empty()); + assert!(plan.recursive_roots.is_empty()); + assert_eq!( + plan.refresh_dirs, + expected_missing_vendor_refresh_seeds(project) + ); + for name in VENDOR_CONFIG_ROOT_NAMES { + let root = project.join(name); + assert!( + plan.refresh_dirs + .contains(&(root.clone(), RecursiveMode::NonRecursive)), + "missing root seed for {name}" + ); + for (dir, mode) in vendor_skill_refresh_dirs(&root) { + assert!( + plan.refresh_dirs.contains(&(dir, mode)), + "missing subdir seed under {name}" + ); + } + } + } + + #[test] + fn plan_skills_watch_targets_does_not_double_seed_present_vendor_roots() { + let tmp = TempDir::new().unwrap(); + let project = tmp.path(); + let grok_home = project.join("home-grok"); + let present: Vec = VENDOR_CONFIG_ROOT_NAMES + .iter() + .map(|name| project.join(name)) + .collect(); + for root in &present { + fs::create_dir_all(root).unwrap(); + } + + let plan = plan_skills_watch_targets(&present, &grok_home, Some(project)); + + assert_eq!(plan.vendor_roots, present); + assert!(plan.project_parent_watch.is_none()); + + let mut expected = Vec::new(); + for root in &present { + expected.extend(vendor_skill_refresh_dirs(root)); + } + assert_eq!(plan.refresh_dirs, expected); + for root in &present { + assert!( + !plan + .refresh_dirs + .contains(&(root.clone(), RecursiveMode::NonRecursive)), + "present vendor root {} must not be refresh-seeded", + root.display() + ); + assert_eq!( + plan.refresh_dirs + .iter() + .filter(|(dir, _)| dir == root || dir.starts_with(root)) + .count(), + vendor_skill_refresh_dirs(root).len() + ); + } + } + + #[test] + fn plan_skills_watch_targets_partial_vendors_seed_only_missing() { + let tmp = TempDir::new().unwrap(); + let project = tmp.path(); + let grok_home = project.join("home-grok"); + let project_claude = project.join(".claude"); + fs::create_dir_all(&project_claude).unwrap(); + + let plan = plan_skills_watch_targets( + std::slice::from_ref(&project_claude), + &grok_home, + Some(project), + ); + + assert_eq!(plan.vendor_roots, vec![project_claude.clone()]); + assert_eq!(plan.project_parent_watch.as_deref(), Some(project)); + + let mut expected = vendor_skill_refresh_dirs(&project_claude).to_vec(); + for name in [".grok", ".agents", ".cursor"] { + let root = project.join(name); + expected.push((root.clone(), RecursiveMode::NonRecursive)); + expected.extend(vendor_skill_refresh_dirs(&root)); + } + assert_eq!(plan.refresh_dirs, expected); + assert!( + !plan + .refresh_dirs + .contains(&(project_claude.clone(), RecursiveMode::NonRecursive)) + ); + } + + #[test] + fn plan_skills_watch_targets_parent_watches_project_when_grok_present_siblings_missing() { + let tmp = TempDir::new().unwrap(); + let project = tmp.path(); + let grok_home = project.join("home-grok"); + let project_grok = project.join(".grok"); + fs::create_dir_all(&project_grok).unwrap(); + + let plan = plan_skills_watch_targets( + std::slice::from_ref(&project_grok), + &grok_home, + Some(project), + ); + + assert_eq!(plan.vendor_roots, vec![project_grok.clone()]); + assert_eq!(plan.project_parent_watch.as_deref(), Some(project)); + assert!( + !plan + .refresh_dirs + .contains(&(project_grok.clone(), RecursiveMode::NonRecursive)) + ); + for name in [".agents", ".claude", ".cursor"] { + let root = project.join(name); + assert!( + plan.refresh_dirs + .contains(&(root.clone(), RecursiveMode::NonRecursive)), + "missing root seed for {name}" + ); + for (dir, mode) in vendor_skill_refresh_dirs(&root) { + assert!( + plan.refresh_dirs.contains(&(dir, mode)), + "missing subdir seed under {name}" + ); + } + } } - /// Regression for the ~/.grok inotify-exhaustion / worktree-noise bug: a - /// `SKILL.md` under a sibling subtree (e.g. `~/.grok/worktrees/`) must not - /// drive a reload, while a real `/skills/**/SKILL.md` change still does. #[test] #[cfg(target_os = "linux")] - fn skills_watcher_scopes_to_subdirs_not_dir_root() { + fn watch_skill_subdirs_ignores_worktrees_under_scoped_root() { let tmp = TempDir::new().unwrap(); let global = tmp.path(); - // Real global skill under /skills/. let alpha = global.join("skills").join("alpha"); fs::create_dir_all(&alpha).unwrap(); fs::write(alpha.join("SKILL.md"), "# alpha").unwrap(); @@ -803,25 +1214,78 @@ mod tests { .expect("debouncer should build"); let watched = watch_skill_subdirs(&mut debouncer, global); - assert!(watched >= 1, "should watch the /skills subdir"); + assert!(watched.contains(&global.join("skills"))); wait_ms(150); - while rx.try_recv().is_ok() {} // drain startup noise + while rx.try_recv().is_ok() {} - // Editing a SKILL.md under the unwatched worktrees/ subtree must NOT fire. fs::write(wt_skill.join("SKILL.md"), "# beta v2").unwrap(); wait_ms(250); assert!( rx.try_recv().is_err(), - "changes below an unwatched sibling subtree (worktrees/) must not \ - trigger a skills reload — proves the dir root is non-recursive" + "changes under worktrees/ must not fire under scoped watches" + ); + + fs::write(alpha.join("SKILL.md"), "# alpha v2").unwrap(); + wait_ms(250); + assert!(rx.try_recv().is_ok(), "changes under skills/ must fire"); + } + + #[test] + #[cfg(target_os = "linux")] + fn watch_skill_subdirs_scopes_project_claude_not_worktrees() { + let tmp = TempDir::new().unwrap(); + let project_claude = tmp.path().join(".claude"); + + let alpha = project_claude.join("skills").join("alpha"); + fs::create_dir_all(&alpha).unwrap(); + fs::write(alpha.join("SKILL.md"), "# alpha").unwrap(); + + let wt_skill = project_claude + .join("worktrees") + .join("wt1") + .join("bazel-out") + .join("deep") + .join("SKILL.md"); + fs::create_dir_all(wt_skill.parent().unwrap()).unwrap(); + fs::write(&wt_skill, "# noise").unwrap(); + + assert!(is_vendor_config_root( + &project_claude, + &tmp.path().join(".grok") + )); + + let (tx, mut rx) = mpsc::unbounded_channel(); + let mut debouncer = new_filtered_debouncer( + Duration::from_millis(50), + move |res: DebounceEventResult| { + let Ok(events) = res else { return }; + if events + .iter() + .any(|event| discovery_change_for_path(&event.path).is_some()) + { + let _ = tx.send(()); + } + }, + ) + .expect("debouncer should build"); + + let watched = watch_skill_subdirs(&mut debouncer, &project_claude); + assert!(watched.contains(&project_claude.join("skills"))); + wait_ms(150); + while rx.try_recv().is_ok() {} + + fs::write(&wt_skill, "# noise v2").unwrap(); + wait_ms(250); + assert!( + rx.try_recv().is_err(), + "changes under project .claude/worktrees must not fire" ); - // Editing the real skill under /skills must still fire. fs::write(alpha.join("SKILL.md"), "# alpha v2").unwrap(); wait_ms(250); assert!( rx.try_recv().is_ok(), - "changes under /skills must trigger a skills reload" + "changes under project .claude/skills must fire" ); } @@ -830,8 +1294,19 @@ mod tests { let grok = Path::new("/tmp/project/.grok"); assert_eq!( discovery_change_for_path(grok), - Some(DiscoveryChange::Skills), - "first .grok creation must take the broader skills reload path" + Some(DiscoveryChange::Skills) + ); + assert_eq!( + discovery_change_for_path(Path::new("/tmp/project/.claude")), + Some(DiscoveryChange::Skills) + ); + assert_eq!( + discovery_change_for_path(&grok.join("skills")), + Some(DiscoveryChange::Skills) + ); + assert_eq!( + discovery_change_for_path(&grok.join("commands")), + Some(DiscoveryChange::Skills) ); assert_eq!( discovery_change_for_path(&grok.join("workflows")), @@ -848,6 +1323,7 @@ mod tests { } #[test] + #[cfg(target_os = "linux")] fn refresh_new_discovery_dirs_attaches_first_created_workflows_dir() { let tmp = TempDir::new().unwrap(); let root = tmp.path(); @@ -886,6 +1362,92 @@ mod tests { assert!(watcher.refreshed_dirs.contains(&workflows)); } + #[test] + fn refresh_new_discovery_dirs_attaches_existing_after_mkdir() { + let tmp = TempDir::new().unwrap(); + let root = tmp.path(); + let workflows = root.join("workflows"); + let mut debouncer = new_filtered_debouncer(Duration::from_millis(50), |_| {}).unwrap(); + debouncer + .watcher() + .watch(root, RecursiveMode::NonRecursive) + .unwrap(); + let mut watcher = SkillsFileWatcher { + debouncer, + refresh_dirs: vec![(workflows.clone(), RecursiveMode::NonRecursive)], + refreshed_dirs: HashSet::new(), + }; + assert!(!watcher.refresh_new_discovery_dirs()); + fs::create_dir(&workflows).unwrap(); + assert!(watcher.refresh_new_discovery_dirs()); + assert!(watcher.refreshed_dirs.contains(&workflows)); + assert!(!watcher.refresh_new_discovery_dirs()); + } + + #[test] + #[cfg(target_os = "linux")] + fn refresh_new_discovery_dirs_attaches_skills_and_commands() { + let tmp = TempDir::new().unwrap(); + let vendor = tmp.path().join(".claude"); + fs::create_dir(&vendor).unwrap(); + + let (tx, mut rx) = mpsc::unbounded_channel(); + let mut debouncer = new_filtered_debouncer( + Duration::from_millis(50), + move |result: DebounceEventResult| { + let Ok(events) = result else { return }; + if events + .iter() + .any(|event| discovery_change_for_path(&event.path).is_some()) + { + let _ = tx.send(()); + } + }, + ) + .unwrap(); + debouncer + .watcher() + .watch(&vendor, RecursiveMode::NonRecursive) + .unwrap(); + + let skills = vendor.join("skills"); + let commands = vendor.join("commands"); + let mut watcher = SkillsFileWatcher { + debouncer, + refresh_dirs: vendor_skill_refresh_dirs(&vendor).to_vec(), + refreshed_dirs: HashSet::new(), + }; + + fs::create_dir_all(skills.join("alpha")).unwrap(); + wait_ms(150); + assert!(rx.try_recv().is_ok(), "must see skills/ creation"); + assert!(watcher.refresh_new_discovery_dirs()); + assert!(watcher.refreshed_dirs.contains(&skills)); + while rx.try_recv().is_ok() {} + + fs::write(skills.join("alpha").join("SKILL.md"), "# alpha").unwrap(); + wait_ms(250); + assert!( + rx.try_recv().is_ok(), + "SKILL.md under newly created skills/ must fire" + ); + while rx.try_recv().is_ok() {} + + fs::create_dir(&commands).unwrap(); + wait_ms(150); + assert!(rx.try_recv().is_ok(), "must see commands/ creation"); + assert!(watcher.refresh_new_discovery_dirs()); + assert!(watcher.refreshed_dirs.contains(&commands)); + while rx.try_recv().is_ok() {} + + fs::write(commands.join("foo.md"), "# foo").unwrap(); + wait_ms(250); + assert!( + rx.try_recv().is_ok(), + "command md under newly created commands/ must fire" + ); + } + #[test] #[cfg_attr( target_os = "macos", diff --git a/crates/codegen/xai-grok-shell/src/extensions/feedback.rs b/crates/codegen/xai-grok-shell/src/extensions/feedback.rs index a52aa79..b752620 100644 --- a/crates/codegen/xai-grok-shell/src/extensions/feedback.rs +++ b/crates/codegen/xai-grok-shell/src/extensions/feedback.rs @@ -18,7 +18,7 @@ use crate::agent::MvpAgent; use crate::session::persistence::{LocalFeedbackEntry, UserFeedbackEntry}; use crate::session::{ ClientFeedbackInput, CommentDeleteRequest, CommentDeleteResponse, CommentRequest, - CommentResponse, FeedbackRequestDismiss, FeedbackResponse, SessionCommand, + CommentResponse, FeedbackRequestDismiss, FeedbackResponse, SessionCommand, SideQuestionError, }; use crate::upload::gcs::WithAuth as _; use xai_file_utils::gcs::upload_bytes; @@ -75,7 +75,20 @@ async fn handle_btw(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { Ok(answer) => super::to_ext_response(Ok(serde_json::json!({ "answer": answer, }))), - Err(e) => Err(acp::Error::internal_error().data(e)), + // Model errors take the canonical mapping: overload gets its short + // display copy there, rate limits keep the typed code + upgrade + // copy, auth failures surface as auth_required. + Err(SideQuestionError::Sampling(e)) => { + Err(crate::sampling::error::map_sampling_err_to_acp(e)) + } + // Non-model failures are already readable sentences. Set `message` + // and leave `data` unset — `Display` appends JSON-encoded `data`, + // and `internal_error().data(e)` rendered as `Internal error: "…"`, + // which made capacity failures look like client bugs in the TUI. + Err(e) => Err(acp::Error::new( + acp::ErrorCode::InternalError.into(), + e.to_string(), + )), } } diff --git a/crates/codegen/xai-grok-shell/src/extensions/notification.rs b/crates/codegen/xai-grok-shell/src/extensions/notification.rs index a61e499..861c9f9 100644 --- a/crates/codegen/xai-grok-shell/src/extensions/notification.rs +++ b/crates/codegen/xai-grok-shell/src/extensions/notification.rs @@ -418,6 +418,25 @@ pub struct HookRunEntryDto { pub output: Option, } +/// Why auto-compaction stopped before completing. +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Eq, + serde::Serialize, + serde::Deserialize, + strum::Display, + strum::EnumString, + strum::AsRefStr, +)] +#[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] +pub enum AutoCompactCancelReason { + UserCancelled, +} + #[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq)] #[serde(rename_all = "snake_case", tag = "sessionUpdate")] pub enum SessionUpdate { @@ -483,7 +502,7 @@ pub enum SessionUpdate { /// Auto-compact was cancelled (user pressed Ctrl+C) AutoCompactCancelled { /// Reason for cancellation - reason: String, + reason: AutoCompactCancelReason, }, /// Auto-continue completed after compaction /// This signals the TUI to flush pending agent messages and end the turn @@ -1704,6 +1723,16 @@ mod tests { let update: SessionUpdate = serde_json::from_str(json).unwrap(); assert_eq!(update, SessionUpdate::MemoryFlushStarted); + // AutoCompactCancelled (strenum reason) + let json = r#"{"sessionUpdate": "auto_compact_cancelled", "reason": "user_cancelled"}"#; + let update: SessionUpdate = serde_json::from_str(json).unwrap(); + assert_eq!( + update, + SessionUpdate::AutoCompactCancelled { + reason: AutoCompactCancelReason::UserCancelled, + } + ); + // AutoCompactFailed (struct variant) let json = r#"{"sessionUpdate": "auto_compact_failed", "error": "oom"}"#; let update: SessionUpdate = serde_json::from_str(json).unwrap(); diff --git a/crates/codegen/xai-grok-shell/src/extensions/session_admin.rs b/crates/codegen/xai-grok-shell/src/extensions/session_admin.rs index d2f7644..d96315f 100644 --- a/crates/codegen/xai-grok-shell/src/extensions/session_admin.rs +++ b/crates/codegen/xai-grok-shell/src/extensions/session_admin.rs @@ -7,6 +7,7 @@ //! - `x.ai/session/rename` rename a session locally + remote //! - `x.ai/session/delete` delete a session locally + remote //! - `x.ai/session/update_mcp_servers` mid-session MCP server swap +//! - `x.ai/session/add_local_workspace` mid-session local workspace add-only (chat) //! - `x.ai/session/fork` fork a session into a new one //! - `x.ai/internal/reload_all_mcp_servers` config hot-reload, all sessions //! - `x.ai/internal/reload_project_mcp_servers` config hot-reload, cwd-scoped @@ -39,6 +40,8 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { "x.ai/session/rename" => handle_session_rename(agent, args).await, "x.ai/session/delete" => handle_session_delete(agent, args).await, "x.ai/session/update_mcp_servers" => handle_update_mcp_servers(agent, args).await, + #[cfg(feature = "local-workspace")] + "x.ai/session/add_local_workspace" => handle_add_local_workspace(agent, args).await, "x.ai/session/fork" => handle_session_fork(agent, args).await, "x.ai/internal/reload_all_mcp_servers" => handle_reload_all_mcp_servers(agent).await, "x.ai/internal/reload_project_mcp_servers" => { @@ -371,6 +374,43 @@ async fn handle_update_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequest) -> .map_err(|e| acp::Error::internal_error().data(e.to_string())) } +// session/add_local_workspace (add-only; local-workspace feature) + +#[cfg(feature = "local-workspace")] +async fn handle_add_local_workspace(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { + #[derive(Deserialize)] + #[serde(rename_all = "camelCase")] + struct Params { + session_id: acp::SessionId, + #[serde(default)] + meta: Option, + } + + let params: Params = parse_params(args)?; + let cwd = { + let sessions = agent.sessions.borrow(); + let h = sessions + .get(¶ms.session_id) + .ok_or_else(|| acp::Error::invalid_params().data("unknown session id"))?; + std::path::PathBuf::from(&h.info.cwd) + }; + // Gate on actual chat kind — not `requires_gateway` (true for non-chat + // GatewayAttach; false for unknown ids). + if !agent.is_chat_kind_session(¶ms.session_id) { + return Err(acp::Error::invalid_params().data(serde_json::json!({ + "code": "local_workspace_chat_only", + "message": "x.ai/session/add_local_workspace is only available on chat-kind sessions", + }))); + } + + let result = agent + .add_local_workspace_mid_session(¶ms.session_id, params.meta, &cwd) + .await?; + ExtMethodResult::success(result) + .to_ext_response() + .map_err(|e| acp::Error::internal_error().data(e.to_string())) +} + // internal/reload_skills /// Reload skills for ALL active sessions. Called by the skills file watcher diff --git a/crates/codegen/xai-grok-shell/src/sampling/error.rs b/crates/codegen/xai-grok-shell/src/sampling/error.rs index ea30d51..2922be6 100644 --- a/crates/codegen/xai-grok-shell/src/sampling/error.rs +++ b/crates/codegen/xai-grok-shell/src/sampling/error.rs @@ -96,12 +96,25 @@ fn pushes_consumer_subscription_upsell(detail: &str) -> bool { d.contains("grok.com/supergrok") || d.contains("upgrade to a grok subscription") } +/// User-facing copy for capacity/overload failures (stream `overloaded_error`, +/// HTTP 529, proxy-wrapped 5xx). See [`SamplingError::is_overloaded`]. +pub const OVERLOADED_USER_MESSAGE: &str = "Model is temporarily overloaded. Try again in a moment."; + /// Map a `SamplingError` to an ACP `Error` for client-facing responses. /// This stays in xai-grok-shell because it depends on `agent_client_protocol::Error`. pub fn map_sampling_err_to_acp(err: SamplingError) -> acp::Error { use reqwest::StatusCode; + // Capacity/overload gets the same short copy on every surface. Message + // only, `data` deliberately unset: `Display` appends JSON-encoded `data`, + // and this string is meant for direct display. + if err.is_overloaded() { + return acp::Error::new( + acp::ErrorCode::InternalError.into(), + OVERLOADED_USER_MESSAGE, + ); + } match err { - SamplingError::Auth(msg) => acp::Error::auth_required().data(msg), + SamplingError::Auth { message, .. } => acp::Error::auth_required().data(message), SamplingError::InvalidConfiguration(msg) => acp::Error::invalid_params().data(msg), SamplingError::Http(e) => { acp::Error::internal_error().data(format!("http client init failed: {e}")) @@ -489,6 +502,31 @@ mod tests { ); } + #[test] + fn overload_maps_to_display_message_without_data() { + let err = SamplingError::StreamError { + error_type: "overloaded_error".into(), + message: "Overloaded".into(), + }; + let acp_err = map_sampling_err_to_acp(err); + assert_eq!(acp_err.code, acp::ErrorCode::InternalError); + assert_eq!(acp_err.message, OVERLOADED_USER_MESSAGE); + // Display appends JSON-encoded `data`; direct-display copy must not + // carry any. + assert_eq!(acp_err.data, None); + + let err_529 = SamplingError::Api { + status: StatusCode::from_u16(529).expect("valid status"), + message: "capacity".into(), + model_metadata: None, + retry_after_secs: None, + should_retry: None, + }; + let acp_529 = map_sampling_err_to_acp(err_529); + assert_eq!(acp_529.message, OVERLOADED_USER_MESSAGE); + assert_eq!(acp_529.data, None); + } + #[test] fn rate_limit_error_uses_dedicated_code() { let err = SamplingError::Api { diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session.rs b/crates/codegen/xai-grok-shell/src/session/acp_session.rs index bb2589c..afae071 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session.rs @@ -63,7 +63,6 @@ use std::sync::Arc; use std::sync::OnceLock; use tokio::sync::{Mutex as TokioMutex, mpsc, oneshot}; use tokio::time::{Duration, sleep}; -use tokio_retry::strategy::ExponentialBackoff; use xai_acp_lib::AcpAgentGatewaySender as GatewaySender; use xai_grok_agent::AgentDefinition; use xai_grok_agent::prompt::agents_md::LEGACY_AGENTS_MD_REMINDER_PREFIX; @@ -92,16 +91,21 @@ mod compaction_segments; mod types; pub(crate) use types::*; pub use types::{TodoGateDecision, TodoGateReason}; +#[path = "acp_session_impl/auth_retry.rs"] +mod auth_retry; #[path = "acp_session_impl/goal.rs"] mod goal; -#[path = "acp_session_impl/interjection.rs"] -mod interjection; -#[path = "acp_session_impl/tool_calls.rs"] -mod tool_calls; #[path = "acp_session_impl/turn.rs"] mod turn; #[path = "acp_session_impl/workflow.rs"] mod workflow_run; +pub(crate) use auth_retry::{ + AuthRetryDecision, AuthRetrySchedule, human_duration, pace_uncharged_resubmit, +}; +#[path = "acp_session_impl/interjection.rs"] +mod interjection; +#[path = "acp_session_impl/tool_calls.rs"] +mod tool_calls; pub(crate) use interjection::*; #[path = "acp_session_impl/laziness.rs"] mod laziness; @@ -1809,6 +1813,9 @@ impl Drop for TurnMetrics { #[cfg(test)] #[path = "acp_session_tests/auth_error_no_retry_tests.rs"] mod auth_error_no_retry_tests; +#[cfg(test)] +#[path = "acp_session_tests/turn/auth_retry_budget_tests.rs"] +mod auth_retry_budget_tests; /// Regression coverage for the auto-wake suppression sweep + shutdown /// drain. These exercise the helpers added to fix the trailing /// `` chat history bug. diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/auth_retry.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/auth_retry.rs new file mode 100644 index 0000000..07ff806 --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/auth_retry.rs @@ -0,0 +1,218 @@ +//! Per-turn retry policy for 401s that follow a *successful* auth recovery +//! (fresh token minted, request to be re-sent). + +use tokio_retry::strategy::ExponentialBackoff; +use xai_grok_sampling_types::SentCredential; + +use super::RecoveredStore; +use crate::auth::AuthManager; +use crate::util::dual_clock::DualClock; + +/// Pace an uncharged resubmit: wait (bounded) for a session-token refresh +/// when that is the store the recovery minted into and nothing wire-valid +/// has landed yet; otherwise floor-pace so the runaway guard can never be a +/// burst of back-to-back requests. Auth policy lives here, not in the turn +/// loop. +pub(crate) async fn pace_uncharged_resubmit( + store: RecoveredStore, + auth_manager: Option<&std::sync::Arc>, +) { + match (store, auth_manager) { + (RecoveredStore::SessionToken, Some(am)) if am.current_wire_valid().is_none() => { + am.wait_for_token_refresh(AuthRetrySchedule::UNCHARGED_REFRESH_WAIT) + .await; + } + _ => tokio::time::sleep(AuthRetrySchedule::UNCHARGED_RESUBMIT_FLOOR).await, + } +} + +/// Compact `2h3m` / `4m7s` / `12s` rendering for turn-failure messages. +pub(crate) fn human_duration(d: std::time::Duration) -> String { + let total_secs = d.as_secs(); + if total_secs < 60 { + return format!("{total_secs}s"); + } + let mins = total_secs / 60; + if mins < 60 { + return format!("{mins}m{}s", total_secs % 60); + } + format!("{}h{}m", mins / 60, mins % 60) +} + +/// Decision for one post-recovery 401 (see +/// [`AuthRetrySchedule::on_recovered_401`]). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AuthRetryDecision { + /// No credential was on the wire, so no slot is charged; resubmit after + /// the refresh lands. `resubmit` is the 1-indexed count since the last + /// successful response. + UnchargedResubmit { resubmit: u32 }, + /// Charged one escalating slot: back off `delay`, then resubmit. + Backoff { + attempt: u32, + delay: std::time::Duration, + }, + /// Per-incident budget exhausted by credentialed 401s — fail the turn. + Exhausted, + /// Runaway guard tripped: recovery kept succeeding while the server + /// rejected `rejections` credential-less requests without a single + /// successful response — fail the turn. + RunawayGuard { rejections: u32 }, +} + +/// Escalating retry budget for post-recovery 401s. The budget is +/// per-incident (successes and suspend boundaries reset it, the latter +/// capped) and only credentialed rejections charge it; the doc on each +/// method carries its own invariant. One thing no reader can derive from +/// the code: +/// +/// **Delays must be 1s/2s/4s.** `ExponentialBackoff::from_millis(base)` +/// raises `base` to the attempt number, so the base must stay small: +/// `from_millis(1000)` yields 1000ⁿ ms = 1s → 16m40s → 11.57 days of +/// silent hang (a past field incident). `from_millis(2).factor(500)` +/// yields 1s, 2s, 4s. +pub(crate) struct AuthRetrySchedule { + delays: std::iter::Take, + /// Slots charged this incident. + attempt: u32, + /// 401s seen this incident, total and the subset that provably carried + /// a credential. Feeds the exhaustion message so "real credential + /// rejected" and "budget exhausted" cannot be conflated. + incident_rejections: u32, + incident_authenticated: u32, + /// Stamped by the incident's first charged 401; cleared by resets. + incident_started: Option, + /// Uncharged fail-closed rejections since the last successful response + /// (survives suspend resets). + uncharged_resubmits: u32, + /// Suspend-triggered resets since the last successful response. + suspend_resets: u32, +} + +impl AuthRetrySchedule { + /// Consecutive credentialed post-recovery 401s tolerated per incident + /// before the turn fails. + pub(crate) const MAX_RETRIES: u32 = 3; + /// Runaway guard: uncharged (no-credential) rejections tolerated + /// without an intervening successful response (~one per 16-minute + /// sleep cycle ⇒ >13 h lid-closed survival). + pub(crate) const MAX_UNCHARGED_RESUBMITS: u32 = 50; + /// Suspend resets tolerated without an intervening successful response + /// (~8 sleep cycles of a continuously failing incident) before the + /// budget stops resetting and is allowed to exhaust. + pub(crate) const MAX_SUSPEND_RESETS: u32 = 8; + /// Bounded wait for the proactive refresh / wake nudge to land a + /// wire-valid token before an uncharged resubmit. + const UNCHARGED_REFRESH_WAIT: std::time::Duration = std::time::Duration::from_secs(15); + /// Floor pacing for uncharged resubmits with no refresh to wait on. + const UNCHARGED_RESUBMIT_FLOOR: std::time::Duration = std::time::Duration::from_secs(1); + /// Wall-vs-monotonic drift beyond which the machine must have slept: + /// well below a real sleep cycle (minutes), well above NTP step jitter. + const SUSPEND_DRIFT_MIN: std::time::Duration = std::time::Duration::from_secs(30); + + pub(crate) fn new() -> Self { + Self { + delays: ExponentialBackoff::from_millis(2) + .factor(500) + .max_delay(std::time::Duration::from_secs(10)) + .take(Self::MAX_RETRIES as usize), + attempt: 0, + incident_rejections: 0, + incident_authenticated: 0, + incident_started: None, + uncharged_resubmits: 0, + suspend_resets: 0, + } + } + + /// Decision for one post-recovery 401. Charges a slot only when the + /// rejected request carried a credential (or its provenance is unknown + /// — fail closed toward terminating). + pub(crate) fn on_recovered_401(&mut self, credential: SentCredential) -> AuthRetryDecision { + self.on_recovered_401_at(credential, DualClock::now()) + } + + /// Clock-injected twin of [`Self::on_recovered_401`] for tests. + fn on_recovered_401_at( + &mut self, + credential: SentCredential, + now: DualClock, + ) -> AuthRetryDecision { + if credential.is_missing() { + self.uncharged_resubmits += 1; + if self.uncharged_resubmits > Self::MAX_UNCHARGED_RESUBMITS { + return AuthRetryDecision::RunawayGuard { + rejections: self.uncharged_resubmits, + }; + } + return AuthRetryDecision::UnchargedResubmit { + resubmit: self.uncharged_resubmits, + }; + } + self.incident_started.get_or_insert(now); + self.incident_rejections += 1; + if credential == SentCredential::Sent { + self.incident_authenticated += 1; + } + match self.delays.next() { + Some(delay) => { + self.attempt += 1; + AuthRetryDecision::Backoff { + attempt: self.attempt, + delay, + } + } + None => AuthRetryDecision::Exhausted, + } + } + + /// Close the open incident if it spans a suspend (wall elapsed outgrew + /// monotonic elapsed by [`Self::SUSPEND_DRIFT_MIN`]): separate wakes are + /// independent 401 events. Capped at [`Self::MAX_SUSPEND_RESETS`] per + /// success-free stretch so a fault that persists across wakes exhausts + /// instead of retrying forever. Returns whether a reset happened. + pub(crate) fn reset_if_incident_spans_suspend(&mut self) -> bool { + self.reset_if_incident_spans_suspend_at(DualClock::now()) + } + + /// Clock-injected twin of [`Self::reset_if_incident_spans_suspend`]. + fn reset_if_incident_spans_suspend_at(&mut self, now: DualClock) -> bool { + let Some(started) = self.incident_started else { + return false; + }; + if self.suspend_resets >= Self::MAX_SUSPEND_RESETS { + return false; + } + let (awake, total) = started.elapsed_between(now); + if total.saturating_sub(awake) < Self::SUSPEND_DRIFT_MIN { + return false; + } + let (uncharged, resets) = (self.uncharged_resubmits, self.suspend_resets); + *self = Self::new(); + self.uncharged_resubmits = uncharged; + self.suspend_resets = resets + 1; + true + } + + /// A successful model response ends every open failure narrative: + /// restart the escalating schedule and clear the success-free-stretch + /// counters (uncharged rejections, suspend resets). + pub(crate) fn reset_on_success(&mut self) { + *self = Self::new(); + } + + /// `(rejections, authenticated)` seen this incident, for the exhaustion + /// message. + pub(crate) fn incident_counts(&self) -> (u32, u32) { + (self.incident_rejections, self.incident_authenticated) + } + + /// Uncharged fail-closed rejections since the last successful response. + pub(crate) fn uncharged_rejections(&self) -> u32 { + self.uncharged_resubmits + } +} + +#[cfg(test)] +#[path = "auth_retry_tests.rs"] +mod tests; diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/auth_retry_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/auth_retry_tests.rs new file mode 100644 index 0000000..6bbf46f --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/auth_retry_tests.rs @@ -0,0 +1,197 @@ +use std::time::Duration; + +use xai_grok_sampling_types::SentCredential; + +use super::{AuthRetryDecision, AuthRetrySchedule}; +use crate::util::dual_clock::DualClock; + +/// `now` shifted `wall_ahead` on the wall clock only — the signature a +/// suspend leaves behind (monotonic pauses, wall keeps advancing). +fn after_suspend(base: DualClock, wall_ahead: Duration) -> DualClock { + DualClock { + mono: base.mono, + wall: base.wall + wall_ahead, + } +} + +/// Pins the exact schedule. Guards against the `from_millis(1000)` footgun +/// (baseⁿ semantics), which produced field sleeps of 1s, 16m40s, and 11.57 +/// days. +#[test] +fn schedule_is_one_two_four_seconds_then_exhausted() { + let mut schedule = AuthRetrySchedule::new(); + let steps: Vec<_> = (0..3) + .map(|_| schedule.on_recovered_401(SentCredential::Sent)) + .collect(); + assert_eq!( + steps, + vec![ + AuthRetryDecision::Backoff { + attempt: 1, + delay: Duration::from_secs(1) + }, + AuthRetryDecision::Backoff { + attempt: 2, + delay: Duration::from_secs(2) + }, + AuthRetryDecision::Backoff { + attempt: 3, + delay: Duration::from_secs(4) + }, + ], + ); + assert_eq!( + schedule.on_recovered_401(SentCredential::Sent), + AuthRetryDecision::Exhausted, + ); + assert_eq!(schedule.incident_counts(), (4, 4)); +} + +/// Unknown provenance charges like an authenticated 401 (fail closed toward +/// terminating) but is not reported as a proven credential rejection. +#[test] +fn unknown_credential_charges_but_is_not_counted_authenticated() { + let mut schedule = AuthRetrySchedule::new(); + assert_eq!( + schedule.on_recovered_401(SentCredential::Unknown), + AuthRetryDecision::Backoff { + attempt: 1, + delay: Duration::from_secs(1) + }, + ); + assert_eq!(schedule.incident_counts(), (1, 0)); +} + +/// The overnight-failure regression: a credential-less 401 never consumes a +/// budget slot; only the runaway guard bounds it. +#[test] +fn missing_credential_never_charges_until_runaway_guard() { + let mut schedule = AuthRetrySchedule::new(); + for i in 1..=AuthRetrySchedule::MAX_UNCHARGED_RESUBMITS { + assert_eq!( + schedule.on_recovered_401(SentCredential::Missing), + AuthRetryDecision::UnchargedResubmit { resubmit: i }, + ); + } + assert_eq!( + schedule.on_recovered_401(SentCredential::Missing), + AuthRetryDecision::RunawayGuard { + rejections: AuthRetrySchedule::MAX_UNCHARGED_RESUBMITS + 1 + }, + ); + assert_eq!( + schedule.on_recovered_401(SentCredential::Sent), + AuthRetryDecision::Backoff { + attempt: 1, + delay: Duration::from_secs(1) + }, + "the credentialed budget must be untouched throughout" + ); +} + +/// A success ends every open failure narrative: the escalating delays, the +/// attempt numbering, and the runaway counter all restart (a 200 disproves +/// the runaway premise, so a productive multi-day turn can never accumulate +/// into the guard). +#[test] +fn success_resets_budget_and_uncharged_counter() { + let mut schedule = AuthRetrySchedule::new(); + schedule.on_recovered_401(SentCredential::Sent); + schedule.on_recovered_401(SentCredential::Sent); + for _ in 0..AuthRetrySchedule::MAX_UNCHARGED_RESUBMITS { + schedule.on_recovered_401(SentCredential::Missing); + } + schedule.reset_on_success(); + assert_eq!( + schedule.on_recovered_401(SentCredential::Missing), + AuthRetryDecision::UnchargedResubmit { resubmit: 1 }, + ); + assert_eq!( + schedule.on_recovered_401(SentCredential::Sent), + AuthRetryDecision::Backoff { + attempt: 1, + delay: Duration::from_secs(1) + }, + ); +} + +/// The uncharged counter survives a suspend reset (the guard spans sleep +/// cycles — that is its point) while the charged budget restarts. +#[test] +fn suspend_reset_preserves_uncharged_counter() { + let mut schedule = AuthRetrySchedule::new(); + let start = DualClock::now(); + schedule.on_recovered_401_at(SentCredential::Missing, start); + schedule.on_recovered_401_at(SentCredential::Missing, start); + schedule.on_recovered_401_at(SentCredential::Sent, start); + + let woke = after_suspend(start, Duration::from_secs(16 * 60)); + assert!(schedule.reset_if_incident_spans_suspend_at(woke)); + assert_eq!( + schedule.on_recovered_401_at(SentCredential::Missing, woke), + AuthRetryDecision::UnchargedResubmit { resubmit: 3 }, + ); + assert_eq!( + schedule.on_recovered_401_at(SentCredential::Sent, woke), + AuthRetryDecision::Backoff { + attempt: 1, + delay: Duration::from_secs(1) + }, + "post-suspend 401 starts a fresh incident instead of exhausting" + ); +} + +/// Suspend resets are capped per success-free stretch: a fault that +/// persists across wakes must eventually exhaust instead of retrying +/// forever. A success re-arms the cap. +#[test] +fn suspend_resets_cap_without_success_and_rearm_on_success() { + let mut schedule = AuthRetrySchedule::new(); + let mut now = DualClock::now(); + for _ in 0..AuthRetrySchedule::MAX_SUSPEND_RESETS { + schedule.on_recovered_401_at(SentCredential::Sent, now); + now = after_suspend(now, Duration::from_secs(16 * 60)); + assert!(schedule.reset_if_incident_spans_suspend_at(now)); + } + schedule.on_recovered_401_at(SentCredential::Sent, now); + now = after_suspend(now, Duration::from_secs(16 * 60)); + assert!( + !schedule.reset_if_incident_spans_suspend_at(now), + "reset {} must be refused: the budget is now allowed to exhaust", + AuthRetrySchedule::MAX_SUSPEND_RESETS + 1 + ); + + schedule.reset_on_success(); + schedule.on_recovered_401_at(SentCredential::Sent, now); + now = after_suspend(now, Duration::from_secs(16 * 60)); + assert!( + schedule.reset_if_incident_spans_suspend_at(now), + "a success re-arms the suspend-reset cap" + ); +} + +/// No suspend, no reset: sub-threshold wall drift (NTP jitter) and a +/// schedule with no open incident are both no-ops. +#[test] +fn suspend_reset_requires_open_incident_and_real_drift() { + let mut schedule = AuthRetrySchedule::new(); + let start = DualClock::now(); + assert!( + !schedule + .reset_if_incident_spans_suspend_at(after_suspend(start, Duration::from_secs(3600))), + "no open incident: nothing to reset" + ); + schedule.on_recovered_401_at(SentCredential::Sent, start); + assert!( + !schedule.reset_if_incident_spans_suspend_at(after_suspend(start, Duration::from_secs(5))), + "5s wall drift is NTP-jitter territory, not a suspend" + ); + assert_eq!( + schedule.on_recovered_401_at(SentCredential::Sent, start), + AuthRetryDecision::Backoff { + attempt: 2, + delay: Duration::from_secs(2) + }, + "the failed reset checks must not charge the budget" + ); +} diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/recap.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/recap.rs index cae2227..d998765 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/recap.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/recap.rs @@ -4,6 +4,37 @@ use super::*; use crate::remote::DEFAULT_CONTEXT_WINDOW; +use crate::session::SideQuestionError; +use xai_grok_sampling_types::SamplingError; + +/// Retry policy for the one-shot `/btw` model call: 3 attempts total +/// (1 try + 2 retries), 500ms → 1s jittered backoff. Deliberately short — +/// nothing like the sampler actor's budget — so a fleet-wide capacity event +/// can't multiply side-question traffic into a retry storm. +fn side_question_retry_policy() -> backon::ExponentialBuilder { + backon::ExponentialBuilder::default() + .with_max_times(2) + .with_min_delay(std::time::Duration::from_millis(500)) + .with_max_delay(std::time::Duration::from_secs(1)) + .with_jitter() +} + +/// Whether a failed `/btw` attempt is worth retrying: overload only (not +/// every retryable 5xx / stream glitch), minus the shared retry vetoes +/// (`x-should-retry: false`, context length — see +/// [`SamplingError::is_retry_vetoed`], also enforced by the sampler actor's +/// `classify_error`). +fn should_retry_side_question(e: &SamplingError) -> bool { + e.is_overloaded() && !e.is_retry_vetoed() +} + +/// Clone the base `/btw` request and stamp a fresh `req_id`, so retried +/// attempts never collide in logs. Everything else is byte-identical. +fn build_side_question_attempt(base: &ConversationRequest) -> ConversationRequest { + let mut request = base.clone(); + request.x_grok_req_id = Some(format!("xai-btw-{}", uuid::Uuid::new_v4())); + request +} impl SessionActor { /// Handle a /btw side question — single-turn model call using the @@ -18,7 +49,10 @@ impl SessionActor { /// /// Generates a unique btw session ID and persists the result to /// `btw_history.jsonl` in the session folder. - pub(super) async fn handle_side_question(&self, question: &str) -> Result { + pub(super) async fn handle_side_question( + &self, + question: &str, + ) -> Result { let btw_session_id = format!("btw-{}", uuid::Uuid::new_v4()); let parent_session_id = self.session_info.id.to_string(); let asked_at = chrono::Utc::now(); @@ -26,7 +60,7 @@ impl SessionActor { let sampling_client = self .prepare_chat_completion(false) .await - .map_err(|e| format!("failed to prepare client: {e}"))?; + .map_err(|e| SideQuestionError::PrepareClient(e.to_string()))?; // Full conversation snapshot including system prompt, tool calls, and results. // Strip reasoning/thinking blocks from assistant items so we don't send @@ -86,7 +120,7 @@ impl SessionActor { .map(|c| c.model) .unwrap_or_default(); - let persist = |answer: String, success: bool, error: Option| { + let persist = |answer: String, success: bool, error: Option, attempts: u32| { let _ = self.notifications.persistence_tx.send(PersistenceMsg::Btw( crate::session::persistence::BtwEntry { btw_session_id: btw_session_id.clone(), @@ -97,6 +131,7 @@ impl SessionActor { model: model.clone(), success, error, + attempts, }, )); }; @@ -105,34 +140,56 @@ impl SessionActor { // `thinking` config via request_defaults for thinking-enabled models, // Anthropic requires temperature == 1 when thinking is enabled. // Leaving it None lets the provider defaults apply correctly. - let request = ConversationRequest { + // + // Built once; each attempt clones it and stamps a fresh req_id (the + // per-attempt clone is the cost of the owned-request API — retries + // are rare, so the success path pays exactly one clone). + let base_request = ConversationRequest { items, tools: tool_specs, model: Some(model.clone()), temperature: None, x_grok_conv_id: Some(btw_session_id.clone()), - x_grok_req_id: Some(format!("xai-btw-{}", uuid::Uuid::new_v4())), x_grok_session_id: Some(parent_session_id.clone()), x_grok_agent_id: Some(xai_grok_telemetry::id::agent_id()), ..Default::default() }; - let response = sampling_client - .conversation_collect(request) - .await - .map_err(|e| { - let msg = format!("side question model call failed: {e}"); - persist(String::new(), false, Some(msg.clone())); - msg - })?; - let content = response.assistant_text(); + // conversation_collect is one-shot (no sampler-actor retry); /btw adds + // its own bounded overload-only retry (policy + predicate above). + use backon::Retryable as _; + let attempts = std::cell::Cell::new(1u32); + let result = + (|| sampling_client.conversation_collect(build_side_question_attempt(&base_request))) + .retry(side_question_retry_policy()) + .when(should_retry_side_question) + .notify(|e: &SamplingError, backoff: std::time::Duration| { + attempts.set(attempts.get() + 1); + tracing::warn!( + backoff_ms = backoff.as_millis() as u64, + error = %e, + "side question overload; retrying" + ); + }) + .await; - if content.is_empty() { - persist(String::new(), false, Some("No response from model".into())); - return Err("No response from model".to_string()); + match result { + Ok(response) => { + let content = response.assistant_text(); + if content.is_empty() { + let err = SideQuestionError::EmptyResponse; + persist(String::new(), false, Some(err.to_string()), attempts.get()); + return Err(err); + } + persist(content.clone(), true, None, attempts.get()); + Ok(content) + } + Err(e) => { + let err = SideQuestionError::from(e); + persist(String::new(), false, Some(err.to_string()), attempts.get()); + Err(err) + } } - persist(content.clone(), true, None); - Ok(content) } /// Generate a session recap and broadcast it via @@ -210,11 +267,10 @@ impl SessionActor { }; let tag = self.reminder_wrapper_tag(); - // Strip reasoning ONLY on the Anthropic Messages backend (it rejects - // thinking blocks without a `thinking` config). Every other backend - // keeps reasoning verbatim so the prefix matches the last turn and the - // provider's prefix KV cache stays warm. Mirrors compaction's - // `summary_strips_reasoning`. + // Strip reasoning only on the Messages backend (it rejects thinking + // blocks without a `thinking` config). Other backends keep reasoning + // verbatim so the prefix matches the last turn and the prefix KV + // cache stays warm. Mirrors compaction's `summary_strips_reasoning`. let strip_reasoning = sampling_client.api_backend() == crate::sampling::ApiBackend::Messages; @@ -689,3 +745,102 @@ impl SessionActor { suggestion } } + +#[cfg(test)] +mod tests { + use super::*; + + fn api(status: u16, message: &str, should_retry: Option) -> SamplingError { + SamplingError::Api { + status: reqwest::StatusCode::from_u16(status).unwrap(), + message: message.into(), + model_metadata: None, + retry_after_secs: None, + should_retry, + } + } + + #[test] + fn side_question_retries_overload_only() { + // Stream overload and its proxy-wrapped 500 shape retry; so does 529. + assert!(should_retry_side_question(&SamplingError::StreamError { + error_type: "overloaded_error".into(), + message: "Overloaded".into(), + })); + assert!(should_retry_side_question(&api( + 500, + "stream error (overloaded_error): Overloaded", + None + ))); + assert!(should_retry_side_question(&api(529, "capacity", None))); + + // Server veto (`x-should-retry: false`) wins over overload. + assert!(!should_retry_side_question(&api( + 529, + "capacity", + Some(false) + ))); + // Deterministic context-length failures never retry, even on 529. + assert!(!should_retry_side_question(&api( + 529, + "invalid_request_error: prompt is too long: 300000 tokens > 200000 maximum", + None + ))); + // Rate limit and generic 5xx are not overload — no /btw retry. + assert!(!should_retry_side_question(&api(429, "slow down", None))); + assert!(!should_retry_side_question(&api( + 503, + "upstream connect timeout", + None + ))); + } + + /// The wired policy: 3 attempts total, backoff within the configured + /// bounds (500ms + 1s base, jitter adds up to the current delay), and a + /// fresh request id stamped per attempt. + #[tokio::test(start_paused = true)] + async fn side_question_retry_wiring_caps_attempts_and_bounds_backoff() { + use backon::Retryable as _; + + let calls = std::cell::Cell::new(0u32); + let start = tokio::time::Instant::now(); + let result: Result<(), SamplingError> = (|| async { + calls.set(calls.get() + 1); + Err(SamplingError::StreamError { + error_type: "overloaded_error".into(), + message: "Overloaded".into(), + }) + }) + .retry(side_question_retry_policy()) + .when(should_retry_side_question) + .await; + + assert!(result.is_err()); + assert_eq!(calls.get(), 3, "1 try + 2 retries"); + // Base delays 500ms + 1s; jitter adds (0, delay) per sleep. + let elapsed = start.elapsed(); + assert!( + elapsed >= std::time::Duration::from_millis(1_500), + "elapsed {elapsed:?} below minimum backoff" + ); + assert!( + elapsed <= std::time::Duration::from_millis(3_100), + "elapsed {elapsed:?} above maximum backoff" + ); + } + + #[test] + fn side_question_attempts_get_fresh_request_ids() { + let base = ConversationRequest { + x_grok_conv_id: Some("btw-test".into()), + ..Default::default() + }; + let a = build_side_question_attempt(&base); + let b = build_side_question_attempt(&base); + let (a_id, b_id) = (a.x_grok_req_id.unwrap(), b.x_grok_req_id.unwrap()); + assert!(a_id.starts_with("xai-btw-")); + assert_ne!(a_id, b_id, "each attempt must get a fresh req_id"); + // Everything except the request id is byte-identical to the base. + assert_eq!(a.x_grok_conv_id, base.x_grok_conv_id); + } +} diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/sampler_turn.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/sampler_turn.rs index fb84ee0..5806afb 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/sampler_turn.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/sampler_turn.rs @@ -925,7 +925,10 @@ impl SessionActor { "auth recovery: sampler 401, devbox re-mint, retrying" ); self.prepare_sampler_for_turn().await; - return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit); + return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { + credential: error.credential, + store: RecoveredStore::SessionToken, + }); } Err(e) => { tracing::warn!( @@ -953,7 +956,10 @@ impl SessionActor { None, ); self.prepare_sampler_for_turn().await; - return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit); + return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { + credential: error.credential, + store: RecoveredStore::SessionToken, + }); } tracing::warn!(session_id = %self.session_info.id.0, "auth recovery: sampler 401, refresh failed"); xai_grok_telemetry::unified_log::warn( @@ -966,7 +972,10 @@ impl SessionActor { && self.try_provider_401_recovery(provider).await { self.prepare_sampler_for_turn().await; - return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit); + return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { + credential: error.credential, + store: RecoveredStore::AuthProvider, + }); } if matches!(error.kind, SamplingErrorKind::IdleTimeout) { self.signals_handle().record_idle_timeout(); @@ -1172,8 +1181,8 @@ impl SessionActor { SamplerFailureRecovery::CompactAndResubmit => { Ok(SamplerTurnOutcome::CompactAndResubmit) } - SamplerFailureRecovery::RefreshAuthAndResubmit => { - Ok(SamplerTurnOutcome::RefreshAuthAndResubmit) + SamplerFailureRecovery::RefreshAuthAndResubmit { credential, store } => { + Ok(SamplerTurnOutcome::RefreshAuthAndResubmit { credential, store }) } } } diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/spawn.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/spawn.rs index 6a0a078..3c8da4b 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/spawn.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/spawn.rs @@ -1593,6 +1593,7 @@ pub(crate) async fn spawn_session_actor( tool_choice: compaction_tool_choice, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: super::memory_state::SessionMemory { flush_config: memory_config.as_ref().map_or_else( diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/tasks_cancel.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/tasks_cancel.rs index 4d7441a..10a03bd 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/tasks_cancel.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/tasks_cancel.rs @@ -248,6 +248,9 @@ impl SessionActor { trigger: Option, ) { let suppress_task_wakes = trigger.as_deref() == Some("ctrl_c"); + // Abort in-flight `/compact` or auto-compact generation (stream select + + // pre-replace guard). Safe when no compact is running. + self.compaction.cancel.request_cancel(); if suppress_task_wakes { if let Some(gate) = &self.tool_context.task_wake_suppressed { gate.set(true); diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/turn.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/turn.rs index a1cfd63..ffabddb 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/turn.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/turn.rs @@ -1,6 +1,7 @@ //! Turn-execution concern for `SessionActor` (`handle_prompt`, turn-end, //! sampling loop). use super::*; +use crate::util::dual_clock::DualClock; use xai_grok_tools::implementations::grok_build::LoopFireMode; /// Synthetic tool the model calls to return its schema-constrained final answer /// on backends that can't constrain output natively (Messages API). Intercepted @@ -1887,6 +1888,7 @@ impl SessionActor { json_schema: Option, ) -> Result { let conv_turn_start = std::time::Instant::now(); + let conv_turn_clock = DualClock::now(); self.maybe_refresh_model_metadata_on_resume().await; self.maybe_compact_on_model_switch().await?; self.chat_state_handle @@ -2193,50 +2195,140 @@ impl SessionActor { return Err(error); } Ok(SamplerTurnOutcome::CompactAndResubmit) => { - auth_retry_schedule.reset(); + auth_retry_schedule.reset_on_success(); continue; } - Ok(SamplerTurnOutcome::RefreshAuthAndResubmit) => { - if let Some((attempt, delay)) = auth_retry_schedule.next_delay() { - let delay_ms = delay.as_millis() as u64; - tracing::warn!( - attempt, - delay_ms, - "auth 401 retry: backing off before resubmit" - ); - xai_grok_telemetry::unified_log::warn( - "shell.turn.auth_retry_backoff", + Ok(SamplerTurnOutcome::RefreshAuthAndResubmit { credential, store }) => { + if auth_retry_schedule.reset_if_incident_spans_suspend() { + tracing::info!("auth 401 retry: incident spanned a suspend; budget reset"); + xai_grok_telemetry::unified_log::info( + "shell.turn.auth_retry_reset_after_suspend", Some(self.session_info.id.0.as_ref()), - Some(serde_json::json!({ - "loop_index": loop_index, - "attempt": attempt, - "max_retries": AuthRetrySchedule::MAX_RETRIES, - "delay_ms": delay_ms, - })), + Some(serde_json::json!({ "loop_index": loop_index })), ); - self.send_xai_notification(XaiSessionUpdate::RetryState( - crate::extensions::notification::RetryState::Retrying { - attempt, - max_retries: AuthRetrySchedule::MAX_RETRIES, - reason: "Re-authenticated after 401; retrying request".to_string(), - }, - )) - .await; - sleep(delay).await; - continue; } - let msg = format!( - "Auth recovery succeeded but inference request was \ - still rejected (401) after {} retries", - AuthRetrySchedule::MAX_RETRIES - ); - tracing::error!(msg); - return Err(acp::Error::internal_error().data( - crate::sampling::error::error_data_with_status(msg, Some(401)), - )); + match auth_retry_schedule.on_recovered_401(credential) { + AuthRetryDecision::UnchargedResubmit { resubmit } => { + tracing::warn!( + resubmit, + "auth 401 retry: no credential was sent; resubmitting uncharged" + ); + xai_grok_telemetry::unified_log::warn( + "shell.turn.auth_resubmit_uncharged", + Some(self.session_info.id.0.as_ref()), + Some(serde_json::json!({ + "loop_index": loop_index, + "resubmit": resubmit, + "max_resubmits": AuthRetrySchedule::MAX_UNCHARGED_RESUBMITS, + })), + ); + self.send_xai_notification(XaiSessionUpdate::RetryState( + crate::extensions::notification::RetryState::Retrying { + attempt: resubmit, + max_retries: AuthRetrySchedule::MAX_UNCHARGED_RESUBMITS, + reason: "Re-authenticated after 401 (request carried no \ + credential); retrying request" + .to_string(), + }, + )) + .await; + pace_uncharged_resubmit(store, self.auth_manager.as_ref()).await; + continue; + } + AuthRetryDecision::Backoff { attempt, delay } => { + let delay_ms = delay.as_millis() as u64; + tracing::warn!( + attempt, + delay_ms, + "auth 401 retry: backing off before resubmit" + ); + xai_grok_telemetry::unified_log::warn( + "shell.turn.auth_retry_backoff", + Some(self.session_info.id.0.as_ref()), + Some(serde_json::json!({ + "loop_index": loop_index, + "attempt": attempt, + "max_retries": AuthRetrySchedule::MAX_RETRIES, + "delay_ms": delay_ms, + })), + ); + self.send_xai_notification(XaiSessionUpdate::RetryState( + crate::extensions::notification::RetryState::Retrying { + attempt, + max_retries: AuthRetrySchedule::MAX_RETRIES, + reason: "Re-authenticated after 401; retrying request" + .to_string(), + }, + )) + .await; + sleep(delay).await; + continue; + } + decision @ (AuthRetryDecision::Exhausted + | AuthRetryDecision::RunawayGuard { .. }) => { + let (awake, wall, suspended) = conv_turn_clock.elapsed_split(); + let duration_note = if suspended >= std::time::Duration::from_secs(1) { + format!( + " Turn ran {} wall-clock, {} of it suspended.", + human_duration(wall), + human_duration(suspended) + ) + } else { + format!(" Turn ran {} wall-clock.", human_duration(wall)) + }; + let (rejections, authenticated) = auth_retry_schedule.incident_counts(); + let uncharged = auth_retry_schedule.uncharged_rejections(); + let msg = match decision { + AuthRetryDecision::RunawayGuard { rejections } => { + format!( + "Auth recovery kept succeeding but {rejections} requests \ + were rejected (401) before a credential could be sent, \ + with no successful response in between; stopping as a \ + runaway guard.{duration_note}" + ) + } + _ if authenticated == rejections => { + format!( + "Auth recovery succeeded but {rejections} authenticated \ + inference requests were still rejected (401); giving up \ + after {} retries.{duration_note}", + AuthRetrySchedule::MAX_RETRIES + ) + } + _ => { + format!( + "Auth retry budget exhausted after {rejections} \ + post-recovery 401s ({authenticated} provably carried a \ + credential).{duration_note}" + ) + } + }; + tracing::error!(msg); + xai_grok_telemetry::unified_log::error( + "shell.turn.auth_retry_exhausted", + Some(self.session_info.id.0.as_ref()), + Some(serde_json::json!({ + "loop_index": loop_index, + "decision": match decision { + AuthRetryDecision::RunawayGuard { .. } => "runaway_guard", + _ => "exhausted", + }, + "rejections": rejections, + "authenticated": authenticated, + "uncharged": uncharged, + "wall_secs": wall.as_secs(), + "awake_secs": awake.as_secs(), + "suspended_secs": suspended.as_secs(), + })), + ); + return Err(acp::Error::internal_error().data( + crate::sampling::error::error_data_with_status(msg, Some(401)), + )); + } + } } }; - auth_retry_schedule.reset(); + auth_retry_schedule.reset_on_success(); let model_elapsed_ms = model_timer.elapsed().as_millis() as u64; let usage = response.usage.as_ref(); let prompt_tokens = usage.map(|u| u.prompt_tokens); @@ -2747,87 +2839,6 @@ mod identical_tool_call_run_tests { assert!(!run.take_nudge()); } } -/// Backoff schedule for resubmits after a *successful* 401 auth recovery -/// (fresh token minted, request to be re-sent). -/// -/// Two hard-won invariants, both regressions from the silent-hang incident -/// where a turn froze 16m40s and then 11.6 days (user-cancelled at 27min): -/// -/// - **Delays must be 1s/2s/4s.** `tokio_retry::ExponentialBackoff:: -/// from_millis(base)` raises `base` to the attempt number, so the base must -/// stay small: `from_millis(1000)` yields 1000ⁿ ms = 1s → 16m40s → 11.57 -/// days. `from_millis(2).factor(500)` yields 2ⁿ × 500ms = 1s, 2s, 4s. -/// - **The schedule is per-incident, not per-turn.** A long turn can span -/// several hourly gateway token rotations; each rotation is an independent -/// 401→refresh→retry event. Without `reset()` after a successful response, -/// the third rotation of one turn would land on the last (largest) delay -/// and the fourth would fail the turn outright. -struct AuthRetrySchedule { - delays: std::iter::Take, - attempt: u32, -} -impl AuthRetrySchedule { - /// Consecutive post-recovery 401s tolerated before the turn fails. - const MAX_RETRIES: u32 = 3; - fn new() -> Self { - Self { - delays: ExponentialBackoff::from_millis(2) - .factor(500) - .max_delay(std::time::Duration::from_secs(10)) - .take(Self::MAX_RETRIES as usize), - attempt: 0, - } - } - /// Next `(attempt_number, delay)` (1-indexed), or `None` once exhausted. - fn next_delay(&mut self) -> Option<(u32, std::time::Duration)> { - let delay = self.delays.next()?; - self.attempt += 1; - Some((self.attempt, delay)) - } - /// A successful model response closes the incident: restart the schedule - /// so the next token rotation starts back at the shortest delay. - fn reset(&mut self) { - *self = Self::new(); - } -} -#[cfg(test)] -mod auth_retry_schedule_tests { - use super::AuthRetrySchedule; - use std::time::Duration; - /// Pins the exact schedule. Guards against the `from_millis(1000)` - /// footgun (baseⁿ semantics): that spelling produced sleeps of 1s, - /// 16m40s, and 11.57 days, observed in the field as a silent - /// ~27-minute hang in `waiting_model` that the user had to cancel. - #[test] - fn schedule_is_one_two_four_seconds_then_exhausted() { - let mut schedule = AuthRetrySchedule::new(); - let steps: Vec<_> = std::iter::from_fn(|| schedule.next_delay()).collect(); - assert_eq!( - steps, - vec![ - (1, Duration::from_secs(1)), - (2, Duration::from_secs(2)), - (3, Duration::from_secs(4)), - ], - ); - assert_eq!( - schedule.next_delay(), - None, - "must exhaust after MAX_RETRIES" - ); - } - /// Each successful response must restart the schedule: hourly token - /// rotations within one long turn are independent incidents, so they - /// must not escalate toward exhaustion (turn failure). - #[test] - fn reset_restarts_delays_and_attempt_numbering() { - let mut schedule = AuthRetrySchedule::new(); - schedule.next_delay(); - schedule.next_delay(); - schedule.reset(); - assert_eq!(schedule.next_delay(), Some((1, Duration::from_secs(1)))); - } -} #[cfg(test)] mod user_echo_broadcast_tests { use super::{UserEchoMode, user_echo_mode}; diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/types.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/types.rs index 5399d92..9679663 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_impl/types.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_impl/types.rs @@ -12,6 +12,20 @@ pub(crate) enum McpReminderMode { Full, } +/// Which credential store a successful 401 recovery minted into. An +/// uncharged resubmit can only usefully wait on the store that recovered: +/// waiting on the session token for a provider-key 401 blocks 15s for a +/// refresh that is irrelevant to the rejected credential. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RecoveredStore { + /// `AuthManager` session token (devbox re-mint, OIDC refresh) — + /// `wait_for_token_refresh` is meaningful. + SessionToken, + /// Auth-provider key minted into chat-state credentials — nothing to + /// wait on in the `AuthManager`; floor-pace instead. + AuthProvider, +} + /// Recovery decision returned by /// `SessionActor::handle_sampling_failure` for the sampler-based /// turn loop. @@ -19,10 +33,15 @@ pub(crate) enum SamplerFailureRecovery { /// Compaction ran. The turn loop should rebuild the request from /// the compacted conversation and resubmit. CompactAndResubmit, - /// Auth 401 recovery succeeded (devbox re-mint, OIDC refresh, or auth - /// provider re-mint). The turn loop should resubmit once with the - /// fresh token. - RefreshAuthAndResubmit, + /// Auth 401 recovery succeeded; the turn loop should resubmit with the + /// fresh token. `credential` is the wire provenance of the rejected + /// request: a 401 for a request that carried no credential at all (a + /// fail-closed send) must not be charged against the per-incident + /// auth-retry budget. + RefreshAuthAndResubmit { + credential: xai_grok_sampling_types::SentCredential, + store: RecoveredStore, + }, } /// Outcome of a single turn attempt via the sampler-based path. @@ -36,8 +55,12 @@ pub(crate) enum SamplerTurnOutcome { Box, ), CompactAndResubmit, - /// Auth recovery succeeded; the outer loop should retry once. - RefreshAuthAndResubmit, + /// Auth recovery succeeded; the outer loop should retry. Mirrors + /// [`SamplerFailureRecovery::RefreshAuthAndResubmit`]. + RefreshAuthAndResubmit { + credential: xai_grok_sampling_types::SentCredential, + store: RecoveredStore, + }, } /// Outcome of `process_conversation_turn`, distinguishing normal completion from cancellation. diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs index 1ac6ab0..e403c57 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs @@ -60,6 +60,7 @@ fn auth_error() -> xai_grok_sampler::SamplingErrorInfo { empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: xai_grok_sampling_types::SentCredential::Unknown, } } @@ -196,7 +197,13 @@ async fn sampler_401_recovery_returns_refresh_and_retry() { let (actor, _rx) = make_actor_with_auth_manager(Some(am)).await; let result = actor.handle_sampling_failure(auth_error()).await; assert!( - matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)), + matches!( + result, + Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { + store: RecoveredStore::SessionToken, + .. + }) + ), "session-based auth with a working refresher must return RefreshAuthAndResubmit" ); assert!(called.load(Ordering::SeqCst), "refresher must be invoked"); @@ -529,6 +536,7 @@ fn model_not_found_error() -> xai_grok_sampler::SamplingErrorInfo { empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: xai_grok_sampling_types::SentCredential::Unknown, } } @@ -596,6 +604,7 @@ fn unauthorized_401_error() -> xai_grok_sampler::SamplingErrorInfo { empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: xai_grok_sampling_types::SentCredential::Unknown, } } @@ -774,7 +783,10 @@ async fn sampler_401_session_method_with_stale_api_key_auth_type_still_recovers( let result = actor.handle_sampling_failure(auth_error()).await; assert!( - matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)), + matches!( + result, + Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { .. }) + ), "session-based method must recover even when auth_type transiently reads ApiKey" ); assert!( @@ -808,7 +820,10 @@ async fn sampler_401_oidc_method_with_stale_api_key_auth_type_still_recovers() { let result = actor.handle_sampling_failure(auth_error()).await; assert!( - matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)), + matches!( + result, + Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { .. }) + ), "oidc method must recover even when auth_type transiently reads ApiKey" ); assert!( @@ -1253,8 +1268,14 @@ async fn sampler_401_on_provider_model_remints_and_resubmits() { let result = actor.handle_sampling_failure(auth_error()).await; assert!( - matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)), - "provider 401 must re-mint and resubmit" + matches!( + result, + Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { + store: RecoveredStore::AuthProvider, + .. + }) + ), + "provider 401 must re-mint and resubmit via the provider store" ); let creds = actor.chat_state_handle.get_credentials().await; assert_eq!( @@ -1289,7 +1310,10 @@ async fn sampler_non_auth_kind_401_on_provider_model_still_recovers() { error.kind = xai_grok_sampler::SamplingErrorKind::Api; let result = actor.handle_sampling_failure(error).await; assert!( - matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)), + matches!( + result, + Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { .. }) + ), "a non-Auth-kind 401 on a provider model must still recover via 4c" ); let creds = actor.chat_state_handle.get_credentials().await; @@ -1321,7 +1345,10 @@ async fn sampler_401_with_no_key_on_provider_model_mints_and_resubmits() { let result = actor.handle_sampling_failure(auth_error()).await; assert!( - matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)), + matches!( + result, + Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { .. }) + ), "an unauthenticated 401 on a provider model must mint and resubmit" ); let creds = actor.chat_state_handle.get_credentials().await; @@ -1364,7 +1391,10 @@ async fn sampler_401_on_provider_model_never_refreshes_session() { let result = actor.handle_sampling_failure(auth_error()).await; assert!( - matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)), + matches!( + result, + Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { .. }) + ), "the provider arm must recover" ); assert!( diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/cancel_running_task_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/cancel_running_task_tests.rs index 3f105b0..0ecb6dd 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/cancel_running_task_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/cancel_running_task_tests.rs @@ -169,6 +169,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() { tool_choice: crate::util::config::CompactionToolChoice::Auto, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: crate::session::memory_state::SessionMemory { flush_config: crate::config::MemoryFlushConfig::default(), @@ -635,6 +636,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history() tool_choice: crate::util::config::CompactionToolChoice::Auto, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: crate::session::memory_state::SessionMemory { flush_config: crate::config::MemoryFlushConfig::default(), @@ -920,6 +922,7 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() { tool_choice: crate::util::config::CompactionToolChoice::Auto, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: crate::session::memory_state::SessionMemory { flush_config: crate::config::MemoryFlushConfig::default(), @@ -2176,6 +2179,7 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() { tool_choice: crate::util::config::CompactionToolChoice::Auto, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: crate::session::memory_state::SessionMemory { flush_config: crate::config::MemoryFlushConfig::default(), diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/idle_resume_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/idle_resume_tests.rs index 688855c..a66b9dc 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/idle_resume_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/idle_resume_tests.rs @@ -195,6 +195,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { tool_choice: crate::util::config::CompactionToolChoice::Auto, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: crate::session::memory_state::SessionMemory { flush_config: crate::config::MemoryFlushConfig::default(), diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs index e00da83..9ccf4de 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs @@ -121,6 +121,7 @@ async fn create_test_actor( tool_choice: crate::util::config::CompactionToolChoice::Auto, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: crate::session::memory_state::SessionMemory { flush_config: crate::config::MemoryFlushConfig::default(), @@ -562,6 +563,7 @@ async fn create_test_actor_with_memory( tool_choice: crate::util::config::CompactionToolChoice::Auto, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: crate::session::memory_state::SessionMemory { flush_config: memory_config @@ -1158,6 +1160,7 @@ fn api_error_with_context_window(context_window: u64) -> xai_grok_sampler::Sampl empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: xai_grok_sampling_types::SentCredential::Unknown, } } /// Primary scenario: remote settings shrinks the context window mid-session. @@ -1345,6 +1348,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { tool_choice: crate::util::config::CompactionToolChoice::Auto, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: crate::session::memory_state::SessionMemory { flush_config: crate::config::MemoryFlushConfig::default(), @@ -1546,6 +1550,7 @@ async fn test_compact_on_error_noop_without_model_metadata() { empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: xai_grok_sampling_types::SentCredential::Unknown, }; assert!(!actor.should_compact_on_error(&err).await); }) diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/memory_config_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/memory_config_tests.rs index 8f0c91c..ff544be 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/memory_config_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/memory_config_tests.rs @@ -171,6 +171,7 @@ async fn create_test_actor_with_memory( tool_choice: crate::util::config::CompactionToolChoice::Auto, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: crate::session::memory_state::SessionMemory { flush_config: memory_config diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs index d4ed2be..2666c8f 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs @@ -123,6 +123,7 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture tool_choice: crate::util::config::CompactionToolChoice::Auto, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: crate::session::memory_state::SessionMemory { flush_config: crate::config::MemoryFlushConfig::default(), @@ -808,6 +809,7 @@ async fn failed_event_preserves_streaming_capture_for_takeout() { empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: xai_grok_sampling_types::SentCredential::Unknown, }, }) .await; @@ -1229,6 +1231,7 @@ async fn reasoning_only_doomloop_turn_captures_every_generation_as_segments() { }), doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: xai_grok_sampling_types::SentCredential::Unknown, }; actor .handle_sampling_event(SamplingEvent::Failed { diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/support.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/support.rs index b910576..26f39c9 100644 --- a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/support.rs +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/support.rs @@ -283,6 +283,7 @@ pub(crate) async fn create_test_actor_ex( tool_choice: crate::util::config::CompactionToolChoice::Auto, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: crate::session::memory_state::SessionMemory { flush_config: crate::config::MemoryFlushConfig::default(), diff --git a/crates/codegen/xai-grok-shell/src/session/acp_session_tests/turn/auth_retry_budget_tests.rs b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/turn/auth_retry_budget_tests.rs new file mode 100644 index 0000000..8d10be9 --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/session/acp_session_tests/turn/auth_retry_budget_tests.rs @@ -0,0 +1,308 @@ +//! Real-turn-loop tests against a mock server that 401s unauthenticated +//! requests and 200s a fresh bearer: a fail-closed (credential-less) 401 +//! must not consume `AuthRetrySchedule` budget — the field failure mode +//! where each sleep cycle burned one slot — while credentialed 401s must +//! still exhaust after `MAX_RETRIES`. + +use super::support::*; +use super::*; +use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::time::Duration; +use xai_grok_test_support::{MockInferenceServer, MockModelEntry}; + +/// The token the mock server accepts and the refresher mints on success. +const FRESH_TOKEN: &str = "refreshed-test-token"; + +/// With `fail_pre_request`, mimics the post-wake sequence: pre-send +/// (`PreRequest`) refreshes fail transiently so the send goes out +/// fail-closed, while the 401-triggered recovery (`ServerRejected`) +/// succeeds and mints [`FRESH_TOKEN`]. Otherwise always succeeds. +struct WakeGapRefresher { + calls: Arc, + fail_pre_request: bool, +} + +#[async_trait::async_trait] +impl crate::auth::refresh::TokenRefresher for WakeGapRefresher { + async fn refresh( + &self, + reason: crate::auth::refresh::RefreshReason, + ) -> crate::auth::refresh::RefreshOutcome { + self.calls.fetch_add(1, Ordering::SeqCst); + if self.fail_pre_request && reason == crate::auth::refresh::RefreshReason::PreRequest { + return crate::auth::refresh::RefreshOutcome::TransientFailure { + message: "simulated post-wake network gap".to_string(), + }; + } + crate::auth::refresh::RefreshOutcome::success(GrokAuth { + key: FRESH_TOKEN.to_string(), + auth_mode: AuthMode::Oidc, + refresh_token: Some("rt-new".into()), + expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), + ..GrokAuth::test_default() + }) + } +} + +/// `(tempdir, manager)` with a hard-expired OIDC token, so the wire-valid +/// resolver has nothing to stamp until the refresher succeeds. The tempdir +/// must outlive the manager (auth.json path). +fn expired_auth_manager( + refresher: Arc, +) -> (tempfile::TempDir, Arc) { + let dir = tempfile::tempdir().expect("tempdir"); + let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default())); + am.hot_swap(GrokAuth { + key: "initial-test-key".into(), + auth_mode: AuthMode::Oidc, + refresh_token: Some("rt".into()), + expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)), + ..GrokAuth::test_default() + }); + am.set_refresher(refresher); + (dir, am) +} + +fn drain_gateway(mut rx: tokio::sync::mpsc::UnboundedReceiver) { + tokio::task::spawn_local(async move { + while let Some(msg) = rx.recv().await { + if let xai_acp_lib::AcpClientMessage::SessionNotification(args) = msg { + let _ = args.response_tx.send(Ok(())); + } + } + }); +} + +fn drain_persistence(mut rx: tokio::sync::mpsc::UnboundedReceiver) { + tokio::task::spawn_local(async move { + while let Some(msg) = rx.recv().await { + if let PersistenceMsg::FlushAndAck { respond_to } = msg { + let _ = respond_to.send(()); + } + } + }); +} + +/// Actor wired for session-token auth against the mock server: real sampler, +/// `cached_token` method, `NotByok` model facts (so the session-token gate is +/// active against the loopback URL), and the supplied auth manager. +async fn session_token_actor( + server: &MockInferenceServer, + auth_manager: Arc, +) -> Arc { + let sampling_cfg = xai_grok_sampler::SamplerConfig { + base_url: server.url(), + model: "test".to_string(), + api_backend: xai_grok_sampler::ApiBackend::Responses, + context_window: 256_000, + max_retries: Some(0), + idle_timeout_secs: Some(30), + ..Default::default() + }; + let (sampler_event_tx, sampler_event_rx) = + tokio::sync::mpsc::unbounded_channel::(); + let sampler_handle = xai_grok_sampler::SamplerActor::spawn( + sampling_cfg, + xai_grok_sampler::RetryPolicy { + max_retries: 0, + rate_limit_retry_threshold: 0, + ..Default::default() + }, + sampler_event_tx, + ); + + let (gateway_tx, gateway_rx) = tokio::sync::mpsc::unbounded_channel(); + drain_gateway(gateway_rx); + let (persistence_tx, persistence_rx) = tokio::sync::mpsc::unbounded_channel(); + drain_persistence(persistence_rx); + + let mut actor = create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await; + actor.sampler_handle = sampler_handle; + actor.auth_manager = Some(auth_manager); + actor.auth_method_id = test_auth_method_id("cached_token"); + + let mut cfg = actor + .chat_state_handle + .get_sampling_config() + .await + .expect("test actor has sampling config"); + cfg.base_url = server.url(); + cfg.api_backend = xai_grok_sampling_types::ApiBackend::Responses; + cfg.model = "test".to_string(); + actor.chat_state_handle.update_sampling_config(cfg); + let mut creds = actor.chat_state_handle.get_credentials().await; + creds.api_key = None; + creds.auth_type = xai_chat_state::AuthType::SessionToken; + actor.chat_state_handle.update_credentials(creds); + + // Definite NotByok: the session-token gate must stay active against the + // loopback mock URL (an `Unknown` would demand a first-party host). + actor + .model_auth_memo + .replace(Some(crate::session::acp_session::ModelAuthMemo { + model_id: "test".to_string(), + facts: crate::agent::config::ModelAuthFacts { + byok: crate::agent::auth_method::ModelByok::NotByok, + auth_scheme: Default::default(), + }, + provider: None, + })); + + actor + .workspace_ops + .bind_local_session( + &actor.session_id_string(), + actor.tool_context.cwd.as_path().to_path_buf(), + actor.tool_context.hunk_tracker_handle.clone(), + actor.agent.borrow().tool_bridge().toolset(), + None, + ) + .expect("bind_local_session"); + + let actor = Arc::new(actor); + { + let drainer = actor.clone(); + let mut sampler_event_rx = sampler_event_rx; + tokio::task::spawn_local(async move { + while let Some(event) = sampler_event_rx.recv().await { + drainer.handle_sampling_event(event).await; + } + }); + } + actor +} + +async fn run_prompt( + actor: &Arc, + prompt_id: &str, +) -> Result { + let prompt_blocks = vec![acp::ContentBlock::Text(acp::TextContent::new( + "hello".to_string(), + ))]; + tokio::time::timeout( + Duration::from_secs(60), + actor.handle_prompt( + prompt_id, + prompt_blocks, + PromptMode::Agent, + None, + None, + None, + None, + true, + None, + None, + None, + ), + ) + .await + .expect("turn must finish within timeout") +} + +/// The wake sequence: the resolver has nothing wire-valid, the send goes +/// out with no `Authorization` header, the server 401s it, recovery lands a +/// fresh token. The turn must survive and resubmit with the fresh bearer. +#[tokio::test(flavor = "current_thread")] +async fn fail_closed_401_is_uncharged_and_turn_survives() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let server = MockInferenceServer::start_with_required_auth( + vec![MockModelEntry::new("test")], + FRESH_TOKEN, + ) + .await + .expect("mock inference server"); + + let calls = Arc::new(AtomicU32::new(0)); + // Pre-send refreshes fail like a post-wake network gap, so the + // first send goes out fail-closed; the 401-recovery refresh + // succeeds. + let refresher = Arc::new(WakeGapRefresher { + calls: calls.clone(), + fail_pre_request: true, + }); + let (_dir, am) = expired_auth_manager(refresher); + let actor = session_token_actor(&server, am).await; + + let outcome = run_prompt(&actor, "auth-retry-budget-fail-closed").await; + assert!( + outcome.is_ok(), + "fail-closed 401 must not fail the turn: {outcome:?}" + ); + + let inference: Vec<_> = server + .requests() + .into_iter() + .filter(|r| r.path.contains("/responses")) + .collect(); + assert!( + inference.len() >= 2, + "expected the fail-closed send plus the resubmit; got {}", + inference.len() + ); + assert_eq!( + inference[0].authorization, None, + "first send must carry no Authorization header" + ); + assert_eq!( + inference.last().unwrap().authorization.as_deref(), + Some(&format!("Bearer {FRESH_TOKEN}") as &str), + "resubmit must carry the freshly refreshed bearer" + ); + assert!( + calls.load(Ordering::SeqCst) >= 2, + "both the failing pre-flight and the recovery refresh must run" + ); + }) + .await; +} + +/// Real credential rejections must still terminate: when every request +/// carries a bearer the server rejects, the escalating budget exhausts after +/// `MAX_RETRIES` and the failure names authenticated rejections — not a +/// generic budget message. `start_paused` auto-advances the backoff ladder. +#[tokio::test(flavor = "current_thread", start_paused = true)] +async fn authenticated_401s_still_exhaust_after_three_retries() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + // The server only accepts a token the refresher never mints, so + // every authenticated send is rejected. + let server = MockInferenceServer::start_with_required_auth( + vec![MockModelEntry::new("test")], + "never-issued-token", + ) + .await + .expect("mock inference server"); + + let refresher = Arc::new(WakeGapRefresher { + calls: Arc::new(AtomicU32::new(0)), + fail_pre_request: false, + }); + let (_dir, am) = expired_auth_manager(refresher); + let actor = session_token_actor(&server, am).await; + + let outcome = run_prompt(&actor, "auth-retry-budget-exhaust").await; + let err = outcome.expect_err("authenticated 401s must exhaust and fail the turn"); + let rendered = serde_json::to_string(&err.data).unwrap_or_default(); + assert!( + rendered.contains("authenticated inference requests were still rejected"), + "exhaustion must name authenticated rejections, got: {rendered}" + ); + + let authenticated = server + .requests() + .into_iter() + .filter(|r| r.path.contains("/responses")) + .filter(|r| r.authorization.as_deref() == Some(&format!("Bearer {FRESH_TOKEN}"))) + .count(); + assert_eq!( + authenticated, 4, + "initial send plus MAX_RETRIES resubmits, all authenticated" + ); + }) + .await; +} diff --git a/crates/codegen/xai-grok-shell/src/session/commands.rs b/crates/codegen/xai-grok-shell/src/session/commands.rs index e709507..6c55d4d 100644 --- a/crates/codegen/xai-grok-shell/src/session/commands.rs +++ b/crates/codegen/xai-grok-shell/src/session/commands.rs @@ -20,6 +20,20 @@ pub struct CancellationContext { /// `None` for graceful in-turn cancels and older clients. pub trigger: Option, } +/// Failure surface of a `/btw` side question. Kept typed until the ACP +/// boundary so `handle_btw` can route model errors through the canonical +/// [`map_sampling_err_to_acp`](crate::sampling::error::map_sampling_err_to_acp) +/// (typed rate-limit / auth codes) instead of a flattened string. +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum SideQuestionError { + #[error("side question model call failed: {0}")] + Sampling(#[from] xai_grok_sampling_types::SamplingError), + #[error("failed to prepare client: {0}")] + PrepareClient(String), + #[error("No response from model")] + EmptyResponse, +} /// Prompt completion kind returned to the ACP layer. #[derive(Debug, Clone)] pub enum PromptCompletionKind { @@ -668,7 +682,7 @@ pub enum SessionCommand { /// tool-free model call, and returns the response text. SideQuestion { question: String, - respond_to: oneshot::Sender>, + respond_to: oneshot::Sender>, }, /// Generate a session recap (a short "where was I" summary) and broadcast /// it to clients via `SessionUpdate::SessionRecap`. diff --git a/crates/codegen/xai-grok-shell/src/session/compaction.rs b/crates/codegen/xai-grok-shell/src/session/compaction.rs index e88cc72..9fb4fb4 100644 --- a/crates/codegen/xai-grok-shell/src/session/compaction.rs +++ b/crates/codegen/xai-grok-shell/src/session/compaction.rs @@ -189,6 +189,7 @@ impl SessionActor { .compaction_policy() .wall_clock_budget_secs; let hosted_tools = self.hosted_tools_for_turn(); + let (cancel, _cancel_scope) = self.compaction.cancel.enter(); match generate_session_compact( history, tools, @@ -199,6 +200,7 @@ impl SessionActor { self.inference_idle_timeout, wall_clock_budget_secs, self.compaction.tool_choice, + &cancel, ) .await { @@ -601,6 +603,7 @@ impl SessionActor { self: &Arc, user_context: Option, ) -> Result<(), acp::Error> { + let (_cancel, _cancel_scope) = self.compaction.cancel.enter(); self.record_compaction_variant(); let total_tokens = self.chat_state_handle.get_total_tokens().await; tracing::Span::current().record("pre_tokens", total_tokens as i64); @@ -638,6 +641,16 @@ impl SessionActor { .await; Ok(()) } + async fn emit_compact_cancelled(&self, auto_trigger: bool) -> Result<(), acp::Error> { + if auto_trigger { + use crate::extensions::notification::SessionUpdate as XaiSessionUpdate; + self.send_xai_notification(XaiSessionUpdate::AutoCompactCancelled { + reason: crate::extensions::notification::AutoCompactCancelReason::UserCancelled, + }) + .await; + } + Err(crate::session::helpers::session_compact::CompactFailure::cancelled_error()) + } /// Suppress AUTO compaction after a deterministic failure. Scope depends on /// the reason (see [`SuppressReason::suppress_state`]): size/schema sticky, /// credit until 200, auth until credentials recover, other clears next turn. @@ -891,6 +904,7 @@ impl SessionActor { auto_continue: Option, trigger: xai_grok_telemetry::events::CompactionTrigger, ) -> Result<(), acp::Error> { + let (cancel, _cancel_scope) = self.compaction.cancel.enter(); let tokens_before = self.chat_state_handle.get_total_tokens().await; tracing::Span::current().record("compaction_tokens_before", tokens_before as i64); self.signals_handle().record_compaction(tokens_before); @@ -1075,6 +1089,7 @@ impl SessionActor { self.inference_idle_timeout, wall_clock_budget_secs, self.compaction.tool_choice, + cancel.clone(), ); let observer = crate::session::helpers::full_replace_compaction::ShellFullReplaceObserver::new( @@ -1135,6 +1150,13 @@ impl SessionActor { deterministic, context_overflow, }) => { + if cancel.is_cancelled() + || message.contains( + crate::session::helpers::session_compact::COMPACT_CANCELLED_MSG, + ) + { + return self.emit_compact_cancelled(auto_trigger).await; + } if context_overflow { let next_stage = match input_stage { InputStage::Verbatim => Some(InputStage::VerbatimFitted), @@ -1575,7 +1597,6 @@ impl SessionActor { let agents_md_reminder = self.agent.borrow().agents_md_user_reminder(); let compaction_context = state_context.for_compaction(); let compaction_state_context: &CompactionStateContext = &compaction_context; - self.persist_compaction_segment(&segment_messages, &generate_session_compact); let transcript_hint = self.transcript_hint(); let summary_count = self .compaction @@ -1620,7 +1641,7 @@ impl SessionActor { user_message_prefix, agents_md_reminder, state_context: &state_context.for_compaction(), - compaction_summary: generate_session_compact, + compaction_summary: generate_session_compact.clone(), system_reminder, summary_before_recent: use_short_prompt, transcript_hint, @@ -1628,8 +1649,6 @@ impl SessionActor { }) }; let prompt_index_at_compaction = self.chat_state_handle.get_prompt_index().await; - self.chat_state_handle - .record_compaction_at(prompt_index_at_compaction); let original_user_info = self .chat_state_handle .get_conversation_item_at(1) @@ -1645,6 +1664,12 @@ impl SessionActor { } _ => None, }); + if cancel.is_cancelled() { + return self.emit_compact_cancelled(auto_trigger).await; + } + self.persist_compaction_segment(&segment_messages, &generate_session_compact); + self.chat_state_handle + .record_compaction_at(prompt_index_at_compaction); self.persist_compaction_checkpoint( &compacted_history, prompt_index_at_compaction, @@ -2008,6 +2033,7 @@ impl SessionActor { trigger_info: AutoCompactTriggerInfo, ) -> Result<(), acp::Error> { use crate::extensions::notification::SessionUpdate as XaiSessionUpdate; + let (_cancel, _cancel_scope) = self.compaction.cancel.enter(); self.record_compaction_variant(); let tokens_before = self.chat_state_handle.get_total_tokens().await; tracing::Span::current().record("pre_tokens", tokens_before as i64); @@ -2058,11 +2084,16 @@ impl SessionActor { let span = tracing::Span::current(); span.record("success", false); span.record("error", e.to_string().as_str()); - if self - .compaction - .auto_compact_suppressed - .load(std::sync::atomic::Ordering::Relaxed) - == SUPPRESS_NONE + let cancelled = self.compaction.cancel.is_cancelled() + || e.data.as_ref().and_then(|d| d.as_str()).is_some_and(|s| { + s.contains(crate::session::helpers::session_compact::COMPACT_CANCELLED_MSG) + }); + if !cancelled + && self + .compaction + .auto_compact_suppressed + .load(std::sync::atomic::Ordering::Relaxed) + == SUPPRESS_NONE { self.send_xai_notification(XaiSessionUpdate::AutoCompactFailed { error: String::new(), @@ -2328,6 +2359,7 @@ mod inline_auto_compact_flow_tests { tool_choice: crate::util::config::CompactionToolChoice::Auto, prefire: crate::session::compaction_config::PrefireState::default(), prefix_released: std::sync::atomic::AtomicBool::new(false), + cancel: Default::default(), }, memory: crate::session::memory_state::SessionMemory { flush_config: crate::config::MemoryFlushConfig::default(), @@ -3683,6 +3715,7 @@ mod inline_auto_compact_flow_tests { empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: xai_grok_sampling_types::SentCredential::Unknown, } } /// Primary scenario: remote settings shrinks the context window mid-session. @@ -3739,6 +3772,7 @@ mod inline_auto_compact_flow_tests { empty_response_context: None, doom_loop_triggers: None, doom_loop_aborted_at_chunk: None, + credential: xai_grok_sampling_types::SentCredential::Unknown, }; assert!(!actor.should_compact_on_error(&err).await); }) diff --git a/crates/codegen/xai-grok-shell/src/session/compaction_config.rs b/crates/codegen/xai-grok-shell/src/session/compaction_config.rs index 497ff1b..2659fd0 100644 --- a/crates/codegen/xai-grok-shell/src/session/compaction_config.rs +++ b/crates/codegen/xai-grok-shell/src/session/compaction_config.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU8; use std::sync::atomic::AtomicU64; +use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; /// Auto-compaction is gated whenever `auto_compact_suppressed` is not [`SUPPRESS_NONE`]. @@ -53,6 +54,58 @@ pub struct AsyncCompactionCache { pub pass1_latency_ms: u64, } +/// Cancel gate for an in-flight compact / prefire sample. +/// +/// Holder count (not a bool): prefire and compact can overlap. The first +/// `enter` installs a token; nested enters reuse it; `in_flight` stays true +/// until the last scope drops. A normal turn stop is a no-op when idle. +#[derive(Default)] +pub struct CompactCancelGate { + token: RefCell, + holders: AtomicUsize, +} + +/// Decrements the holder count when a compact/prefire scope ends. +pub struct CompactCancelScope<'a>(&'a CompactCancelGate); + +impl Drop for CompactCancelScope<'_> { + fn drop(&mut self) { + self.0.end(); + } +} + +impl CompactCancelGate { + /// Start or join a compact/prefire scope. Nested callers share one token, + /// including a token already cancelled by stop, so overlapping prefire + + /// compact both observe the same abort. A later independent enter after + /// holders drain installs a fresh token. + pub fn enter(&self) -> (tokio_util::sync::CancellationToken, CompactCancelScope<'_>) { + let prev = self.holders.fetch_add(1, Ordering::AcqRel); + let token = if prev == 0 { + let token = tokio_util::sync::CancellationToken::new(); + self.token.replace(token.clone()); + token + } else { + self.token.borrow().clone() + }; + (token, CompactCancelScope(self)) + } + + fn end(&self) { + self.holders.fetch_sub(1, Ordering::AcqRel); + } + + pub fn request_cancel(&self) { + if self.holders.load(Ordering::Acquire) > 0 { + self.token.borrow().cancel(); + } + } + + pub fn is_cancelled(&self) -> bool { + self.holders.load(Ordering::Acquire) > 0 && self.token.borrow().is_cancelled() + } +} + /// Prefire two-pass state. `Default` so it drops into existing `CompactionConfig` /// struct literals with a single `prefire: PrefireState::default()` field. /// @@ -152,6 +205,8 @@ pub struct CompactionConfig { pub prefire: PrefireState, /// Sticky once a forked session releases its inherited prefix under compaction pressure (see `run_compact_inner`), so it stops re-pinning it. pub prefix_released: AtomicBool, + /// User/stop cancel for the current compact generation. + pub cancel: CompactCancelGate, } #[cfg(test)] @@ -209,3 +264,62 @@ mod prefire_state_tests { assert!(state.take().is_none()); } } + +#[cfg(test)] +mod compact_cancel_gate_tests { + use super::*; + + #[test] + fn request_cancel_trips_shared_token() { + let gate = CompactCancelGate::default(); + let (token, _scope) = gate.enter(); + assert!(!token.is_cancelled()); + gate.request_cancel(); + assert!(token.is_cancelled()); + assert!(gate.is_cancelled()); + } + + #[test] + fn request_cancel_is_noop_when_idle() { + let gate = CompactCancelGate::default(); + gate.request_cancel(); + let (token, _scope) = gate.enter(); + assert!(!token.is_cancelled()); + assert!(!gate.is_cancelled()); + } + + #[test] + fn nested_enter_keeps_in_flight_after_inner_drop() { + let gate = CompactCancelGate::default(); + let (outer_tok, outer) = gate.enter(); + let (inner_tok, inner) = gate.enter(); + gate.request_cancel(); + assert!(outer_tok.is_cancelled()); + assert!(inner_tok.is_cancelled()); + drop(inner); + assert!(gate.is_cancelled()); + drop(outer); + assert!(!gate.is_cancelled()); + let (next, _scope) = gate.enter(); + assert!(!next.is_cancelled()); + } + + #[test] + fn join_while_cancelled_reuses_cancelled_token() { + let gate = CompactCancelGate::default(); + let (_outer, outer) = gate.enter(); + gate.request_cancel(); + let (joined, joined_scope) = gate.enter(); + assert!( + joined.is_cancelled(), + "nested enter during stop must keep sharing the cancelled token" + ); + drop(joined_scope); + drop(outer); + let (next, _scope) = gate.enter(); + assert!( + !next.is_cancelled(), + "fresh enter after scopes drain must not inherit the prior stop" + ); + } +} diff --git a/crates/codegen/xai-grok-shell/src/session/helpers/full_replace_compaction.rs b/crates/codegen/xai-grok-shell/src/session/helpers/full_replace_compaction.rs index c84ad29..d023f6b 100644 --- a/crates/codegen/xai-grok-shell/src/session/helpers/full_replace_compaction.rs +++ b/crates/codegen/xai-grok-shell/src/session/helpers/full_replace_compaction.rs @@ -72,6 +72,7 @@ pub(crate) struct ShellCompactionSampler { /// reasoning-runaway backstop; `0` disables it. wall_clock_budget_secs: u64, tool_choice: crate::util::config::CompactionToolChoice, + cancel: tokio_util::sync::CancellationToken, /// Full output of the most recent successful sample (for L5 telemetry). last_success: Mutex>, } @@ -89,6 +90,7 @@ impl ShellCompactionSampler { idle_timeout: Duration, wall_clock_budget_secs: u64, tool_choice: crate::util::config::CompactionToolChoice, + cancel: tokio_util::sync::CancellationToken, ) -> Self { Self { use_short_prompt, @@ -101,6 +103,7 @@ impl ShellCompactionSampler { idle_timeout, wall_clock_budget_secs, tool_choice, + cancel, last_success: Mutex::new(None), } } @@ -140,6 +143,7 @@ impl CompactionSampler for ShellCompactionSampler { self.idle_timeout, self.wall_clock_budget_secs, self.tool_choice, + &self.cancel, ) .await { @@ -170,6 +174,7 @@ fn compact_failure_to_sample_error(failure: CompactFailure) -> CompactionSampleE let (deterministic, err) = match failure { CompactFailure::Deterministic(err) => (true, err), CompactFailure::Transient(err) => (false, err), + CompactFailure::Cancelled => (true, CompactFailure::cancelled_error()), }; let message = acp_error_message(&err); if deterministic { diff --git a/crates/codegen/xai-grok-shell/src/session/helpers/session_compact.rs b/crates/codegen/xai-grok-shell/src/session/helpers/session_compact.rs index 9db539c..1ada944 100644 --- a/crates/codegen/xai-grok-shell/src/session/helpers/session_compact.rs +++ b/crates/codegen/xai-grok-shell/src/session/helpers/session_compact.rs @@ -45,6 +45,15 @@ pub(crate) enum CompactFailure { /// Failure may resolve on retry. The caller follows its existing /// N-attempt + backoff loop. Transient(acp::Error), + /// User/stop cancelled the in-flight compact. Do not retry or suppress AUTO. + Cancelled, +} +/// Stable error payload for a user-cancelled compact (pager + retry loop). +pub(crate) const COMPACT_CANCELLED_MSG: &str = "compact cancelled"; +impl CompactFailure { + pub(crate) fn cancelled_error() -> acp::Error { + acp::Error::internal_error().data(COMPACT_CANCELLED_MSG) + } } pub(crate) use xai_grok_sampling_types::is_context_length_error; /// Classify an upstream `SamplingError` for the compaction retry loop. @@ -58,7 +67,7 @@ pub(crate) use xai_grok_sampling_types::is_context_length_error; fn classify_sampling_error(err: SamplingError) -> CompactFailure { let acp_err = acp::Error::internal_error().data(format!("compact failed: {err}")); let deterministic = match &err { - SamplingError::Auth(_) + SamplingError::Auth { .. } | SamplingError::InvalidConfiguration(_) | SamplingError::Serialization(_) | SamplingError::IdleTimeout { .. } => true, @@ -315,14 +324,75 @@ enum StreamStep { Ended, IdleTimeout, } -async fn next_stream_step(stream: &mut S, idle_timeout: std::time::Duration) -> StreamStep +async fn next_stream_step( + stream: &mut S, + idle_timeout: std::time::Duration, + cancel: &tokio_util::sync::CancellationToken, +) -> Result, CompactFailure> where S: futures_util::Stream + Unpin, { - match tokio::time::timeout(idle_timeout, stream.next()).await { - Ok(Some(item)) => StreamStep::Item(item), - Ok(None) => StreamStep::Ended, - Err(_) => StreamStep::IdleTimeout, + tokio::select! { + biased; + _ = cancel.cancelled() => Err(CompactFailure::Cancelled), + step = tokio::time::timeout(idle_timeout, stream.next()) => Ok(match step { + Ok(Some(item)) => StreamStep::Item(item), + Ok(None) => StreamStep::Ended, + Err(_) => StreamStep::IdleTimeout, + }), + } +} +/// Abort `fut` if stop wins while the compact HTTP stream is still opening. +async fn await_unless_cancelled( + cancel: &tokio_util::sync::CancellationToken, + fut: F, +) -> Result +where + F: std::future::Future, +{ + tokio::select! { + biased; + _ = cancel.cancelled() => Err(CompactFailure::Cancelled), + result = fut => Ok(result), + } +} +#[cfg(test)] +mod compact_cancel_await_tests { + use super::*; + use std::time::Duration; + use tokio_util::sync::CancellationToken; + #[tokio::test] + async fn pre_cancelled_token_skips_fut() { + let cancel = CancellationToken::new(); + cancel.cancel(); + let err = await_unless_cancelled(&cancel, async { + panic!("fut must not run when already cancelled"); + }) + .await + .unwrap_err(); + assert!(matches!(err, CompactFailure::Cancelled)); + } + #[tokio::test] + async fn cancel_aborts_pending_open() { + let cancel = CancellationToken::new(); + let cancel2 = cancel.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(20)).await; + cancel2.cancel(); + }); + let started = std::time::Instant::now(); + let err = await_unless_cancelled(&cancel, async { + tokio::time::sleep(Duration::from_secs(30)).await; + 0u8 + }) + .await + .unwrap_err(); + assert!(matches!(err, CompactFailure::Cancelled)); + assert!( + started.elapsed() < Duration::from_secs(2), + "stop must abort stream-open wait, elapsed {:?}", + started.elapsed() + ); } } /// Generates a summary of the conversation for compaction. @@ -355,7 +425,11 @@ pub(crate) async fn generate_session_compact( idle_timeout: std::time::Duration, wall_clock_budget_secs: u64, tool_choice: crate::util::config::CompactionToolChoice, + cancel: &tokio_util::sync::CancellationToken, ) -> Result { + if cancel.is_cancelled() { + return Err(CompactFailure::Cancelled); + } let num_messages = chat_history.len(); let wire_tool_choice = match tool_choice { crate::util::config::CompactionToolChoice::Auto => ToolChoice::auto(), @@ -392,7 +466,8 @@ pub(crate) async fn generate_session_compact( num_messages = num_messages, "Sending compact request (streaming)" ); - let stream_result = client.chat_completion_stream(message).await; + let stream_result = + await_unless_cancelled(cancel, client.chat_completion_stream(message)).await?; let mut stream = match stream_result { Ok((s, _metadata)) => s, Err(e) => return Err(classify_sampling_error(e)), @@ -404,7 +479,9 @@ pub(crate) async fn generate_session_compact( let mut last_progress_at = std::time::Instant::now(); loop { let idle_remaining = idle_timeout.saturating_sub(last_progress_at.elapsed()); - let chunk_result = match next_stream_step(&mut stream, idle_remaining).await { + let chunk_result = match next_stream_step(&mut stream, idle_remaining, cancel) + .await? + { StreamStep::Item(item) => item, StreamStep::Ended => break, StreamStep::IdleTimeout => { @@ -486,7 +563,9 @@ pub(crate) async fn generate_session_compact( x_grok_agent_id: Some(xai_grok_telemetry::id::agent_id()), ..Default::default() }; - let stream_result = client.conversation_stream_responses(request).await; + let stream_result = + await_unless_cancelled(cancel, client.conversation_stream_responses(request)) + .await?; let mut stream = match stream_result { Ok((s, _metadata, _doom_loop)) => s, Err(e) => return Err(classify_sampling_error(e)), @@ -498,7 +577,9 @@ pub(crate) async fn generate_session_compact( let mut last_progress_at = std::time::Instant::now(); loop { let idle_remaining = idle_timeout.saturating_sub(last_progress_at.elapsed()); - let chunk_result = match next_stream_step(&mut stream, idle_remaining).await { + let chunk_result = match next_stream_step(&mut stream, idle_remaining, cancel) + .await? + { StreamStep::Item(item) => item, StreamStep::Ended => break, StreamStep::IdleTimeout => { @@ -611,7 +692,9 @@ pub(crate) async fn generate_session_compact( x_grok_agent_id: Some(xai_grok_telemetry::id::agent_id()), ..Default::default() }; - let stream_result = client.conversation_stream_messages(request).await; + let stream_result = + await_unless_cancelled(cancel, client.conversation_stream_messages(request)) + .await?; let mut stream = match stream_result { Ok((s, _metadata)) => s, Err(e) => return Err(classify_sampling_error(e)), @@ -623,7 +706,9 @@ pub(crate) async fn generate_session_compact( let mut last_progress_at = std::time::Instant::now(); loop { let idle_remaining = idle_timeout.saturating_sub(last_progress_at.elapsed()); - let chunk_result = match next_stream_step(&mut stream, idle_remaining).await { + let chunk_result = match next_stream_step(&mut stream, idle_remaining, cancel) + .await? + { StreamStep::Item(item) => item, StreamStep::Ended => break, StreamStep::IdleTimeout => { @@ -771,9 +856,9 @@ mod classify_tests { } #[test] fn sampling_non_api_variants_classify_correctly() { - assert!(is_det(&classify_sampling_error(SamplingError::Auth( - "expired".into() - )))); + assert!(is_det(&classify_sampling_error( + SamplingError::auth_unknown("expired") + ))); assert!(is_det(&classify_sampling_error( SamplingError::InvalidConfiguration("missing key") ))); @@ -1619,6 +1704,7 @@ mod reasoning_compaction_regression_tests { std::time::Duration::from_secs(30), 0, crate::util::config::CompactionToolChoice::Auto, + &tokio_util::sync::CancellationToken::new(), ) .await .unwrap_or_else(|_| panic!("compaction must succeed")); @@ -1710,6 +1796,7 @@ mod reasoning_compaction_regression_tests { std::time::Duration::from_secs(30), 0, crate::util::config::CompactionToolChoice::Auto, + &tokio_util::sync::CancellationToken::new(), ) .await; let output = result @@ -1772,6 +1859,7 @@ mod reasoning_compaction_regression_tests { std::time::Duration::from_secs(30), 0, crate::util::config::CompactionToolChoice::Auto, + &tokio_util::sync::CancellationToken::new(), ) .await .unwrap_or_else(|_| panic!("compaction with tools must succeed")); @@ -1786,6 +1874,7 @@ mod reasoning_compaction_regression_tests { std::time::Duration::from_secs(30), 0, crate::util::config::CompactionToolChoice::Auto, + &tokio_util::sync::CancellationToken::new(), ) .await .unwrap_or_else(|_| panic!("compaction without tools must succeed")); @@ -1919,6 +2008,7 @@ mod reasoning_compaction_regression_tests { std::time::Duration::from_secs(30), 0, crate::util::config::CompactionToolChoice::Auto, + &tokio_util::sync::CancellationToken::new(), ) .await .unwrap_or_else(|_| panic!("Responses compaction with tools must succeed")); @@ -1933,6 +2023,7 @@ mod reasoning_compaction_regression_tests { std::time::Duration::from_secs(30), 0, crate::util::config::CompactionToolChoice::Auto, + &tokio_util::sync::CancellationToken::new(), ) .await .unwrap_or_else(|_| panic!("Responses compaction without tools must succeed")); @@ -2012,6 +2103,7 @@ mod reasoning_compaction_regression_tests { std::time::Duration::from_millis(150), 0, crate::util::config::CompactionToolChoice::Auto, + &tokio_util::sync::CancellationToken::new(), ) .await; match result { @@ -2026,8 +2118,10 @@ mod reasoning_compaction_regression_tests { "expected an idle-timeout transient failure, got: {data}" ); } - Err(CompactFailure::Deterministic(_)) => { - panic!("a stalled stream must be retryable (Transient), not Deterministic") + Err(CompactFailure::Deterministic(_) | CompactFailure::Cancelled) => { + panic!( + "a stalled stream must be retryable (Transient), not Deterministic/Cancelled" + ) } Ok(_) => panic!("a stalled stream must not produce a summary"), } @@ -2091,6 +2185,7 @@ mod reasoning_compaction_regression_tests { std::time::Duration::from_millis(150), 0, crate::util::config::CompactionToolChoice::Auto, + &tokio_util::sync::CancellationToken::new(), ) .await; match result { @@ -2105,7 +2200,7 @@ mod reasoning_compaction_regression_tests { "expected an idle-timeout transient failure, got: {data}" ); } - Err(CompactFailure::Deterministic(_)) => { + Err(CompactFailure::Deterministic(_) | CompactFailure::Cancelled) => { panic!("a stalled stream must be retryable (Transient), not Deterministic") } Ok(_) => { @@ -2169,6 +2264,7 @@ mod reasoning_compaction_regression_tests { std::time::Duration::from_millis(150), 0, crate::util::config::CompactionToolChoice::Auto, + &tokio_util::sync::CancellationToken::new(), ) .await; match result { @@ -2183,7 +2279,7 @@ mod reasoning_compaction_regression_tests { "expected an idle-timeout transient failure, got: {data}" ); } - Err(CompactFailure::Deterministic(_)) => { + Err(CompactFailure::Deterministic(_) | CompactFailure::Cancelled) => { panic!("a stalled stream must be retryable (Transient), not Deterministic") } Ok(_) => { @@ -2244,6 +2340,7 @@ mod reasoning_compaction_regression_tests { std::time::Duration::from_millis(150), 0, crate::util::config::CompactionToolChoice::Auto, + &tokio_util::sync::CancellationToken::new(), ) .await; match result { diff --git a/crates/codegen/xai-grok-shell/src/session/persistence.rs b/crates/codegen/xai-grok-shell/src/session/persistence.rs index 8e55706..4ecfe44 100644 --- a/crates/codegen/xai-grok-shell/src/session/persistence.rs +++ b/crates/codegen/xai-grok-shell/src/session/persistence.rs @@ -73,6 +73,14 @@ pub struct BtwEntry { /// Error message if failed. #[serde(skip_serializing_if = "Option::is_none")] pub error: Option, + /// Model-call attempts made (1 = no retry). Entries written before this + /// field existed deserialize as 1. + #[serde(default = "default_btw_attempts")] + pub attempts: u32, +} + +fn default_btw_attempts() -> u32 { + 1 } // Local feedback persistence types diff --git a/crates/codegen/xai-grok-shell/src/session/unified_list/mod.rs b/crates/codegen/xai-grok-shell/src/session/unified_list/mod.rs index 3ea7574..3dd44d0 100644 --- a/crates/codegen/xai-grok-shell/src/session/unified_list/mod.rs +++ b/crates/codegen/xai-grok-shell/src/session/unified_list/mod.rs @@ -55,10 +55,18 @@ pub fn conversations_lane_active() -> bool { } /// Parse `x.ai/session/list` params and, under process-wide chat mode, force /// the conversations-only `kind` facet (see [`force_kind_chat`]). +/// +/// Client-sent `kind` of `chat`/`build` is honored only behind +/// `feature = "local-workspace"` (pager welcome Local history). Chat-only +/// Desktop/ACP agents keep the force-rewrite so `kind: ["build"]` cannot +/// surface Build rows. 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); + let honor_client_kind = cfg!(feature = "local-workspace") && client_sent_kind_filter(&req); + if !honor_client_kind { + force_kind_chat(&mut req); + } } Ok(req) } @@ -72,6 +80,23 @@ where CwdScope::WithSiblings }) } +fn client_sent_kind_filter(req: &ListReq) -> bool { + let Some(kind) = req + .meta + .as_ref() + .and_then(|m| m.get("x.ai/facetFilters")) + .and_then(|f| f.get("kind")) + else { + return false; + }; + match kind { + serde_json::Value::Array(arr) if !arr.is_empty() => arr + .iter() + .any(|v| matches!(v.as_str(), Some("chat" | "build"))), + serde_json::Value::String(s) if s == "chat" || s == "build" => true, + _ => false, + } +} #[derive(Debug, Default, Deserialize)] #[serde(rename_all = "camelCase")] pub struct ListReq { @@ -172,6 +197,10 @@ fn value_list(v: &serde_json::Value) -> Vec { } /// Rewrite `req` so the `kind` facet filter is exactly `["chat"]`. /// +/// Used when process chat mode is on **and** the client omitted a recognized +/// `kind` facet (see [`parse_list_req`]). Welcome history sends an explicit +/// `kind` (`chat` / `build`) that must not be rewritten. Other facet filters +/// and `_meta` keys are left untouched. pub fn force_kind_chat(req: &mut ListReq) { force_kind(req, SessionKind::Chat); } @@ -201,7 +230,7 @@ pub async fn build_unified_list( conversations_client: Option<&ConversationsClient>, mut req: ListReq, ) -> UnifiedListResult { - if crate::agent::chat_modes::process_chat_mode_enabled() { + if crate::agent::chat_modes::process_chat_mode_enabled() && !client_sent_kind_filter(&req) { force_kind_chat(&mut req); } let reg = facet_registry(); @@ -936,16 +965,42 @@ mod tests { let _on = xai_grok_test_support::EnvGuard::set(GROK_CHAT_MODE_ENV, "1"); let req = parse_list_req(&raw).expect("parse"); let parsed = ParsedMeta::parse(req.meta.as_ref()); - let expected = "build"; + let expected_build = if cfg!(feature = "local-workspace") { + Some(&vec![serde_json::json!("build")]) + } else { + Some(&vec![serde_json::json!("build")]) + }; assert_eq!( parsed.facet_filters.get(KIND_FACET_KEY), - Some(&vec![serde_json::json!(expected)]) + expected_build, + "client kind=build under process chat mode" ); assert_eq!( parsed.facet_filters.get("starred"), Some(&vec![serde_json::json!(true)]), "other facets pass through" ); + let req = parse_list_req("{}").expect("parse"); + let parsed = ParsedMeta::parse(req.meta.as_ref()); + let expected = None; + assert_eq!( + parsed.facet_filters.get(KIND_FACET_KEY), + expected, + "absent client kind still forces chat under process chat mode" + ); + for bad in [ + serde_json::json!({ "_meta": { "x.ai/facetFilters": { "kind": [] } } }), + serde_json::json!({ "_meta": { "x.ai/facetFilters": { "kind": null } } }), + serde_json::json!({ "_meta": { "x.ai/facetFilters": { "kind": ["other"] } } }), + ] { + let req = parse_list_req(&bad.to_string()).expect("parse"); + let parsed = ParsedMeta::parse(req.meta.as_ref()); + assert_eq!( + parsed.facet_filters.get(KIND_FACET_KEY), + expected, + "empty/null/unknown kind must still force chat: {bad}" + ); + } } } /// Wire pin for the cross-crate `x.ai/partial` envelope the pager parses: diff --git a/crates/codegen/xai-grok-shell/src/terminal/pty_session.rs b/crates/codegen/xai-grok-shell/src/terminal/pty_session.rs index 6ea0633..2bfdd4b 100644 --- a/crates/codegen/xai-grok-shell/src/terminal/pty_session.rs +++ b/crates/codegen/xai-grok-shell/src/terminal/pty_session.rs @@ -5,6 +5,10 @@ //! on its own leaves them running, as a terminal does: their process groups are //! their own and nothing here holds a handle to them. +// A panic here loses a shell: teardown paths run inside `Drop`, where an +// unwind during another unwind aborts the process. Tests panic freely. +#![cfg_attr(not(test), deny(clippy::unwrap_used, clippy::expect_used))] + use std::collections::{HashMap, VecDeque}; use std::io::{Read, Write}; use std::sync::{Arc, LazyLock}; @@ -147,6 +151,53 @@ impl Shell { } } +/// A shell not yet in the registry, where teardown would never find it. +/// Dropping reaps it; [`Self::into_registered`] hands it to the registry +/// instead. +/// +/// Disarming leaves [`Shell::Reaped`] behind rather than an empty slot, so the +/// guard has no state in which its own field is missing. +struct UnregisteredShell(Shell); + +impl UnregisteredShell { + fn new(child: Box) -> Self { + Self(Shell::Running { child, group: None }) + } + + fn pid(&self) -> Option { + self.0.pid() + } + + fn attach_group(&mut self, enrolled: Arc) { + self.0.attach_group(enrolled); + } + + /// Disarms the guard. The caller must reach the registry without awaiting, + /// or it reopens the window this type closes. + fn into_registered(mut self) -> Shell { + self.disarm() + } + + fn disarm(&mut self) -> Shell { + std::mem::replace(&mut self.0, Shell::Reaped(None)) + } +} + +impl Drop for UnregisteredShell { + fn drop(&mut self) { + let mut shell = self.disarm(); + // `reap_now` blocks through its grace waits, so keep it off an async + // thread. Not the runtime's blocking pool though: a task still queued + // there at shutdown is dropped unrun, taking the group with it and + // leaving the scope holding a dead `Weak`. + if tokio::runtime::Handle::try_current().is_ok() { + std::thread::spawn(move || shell.reap_now()); + } else { + shell.reap_now(); + } + } +} + type PtyMap = HashMap>>; static PTY_REGISTRY: LazyLock> = LazyLock::new(|| Mutex::new(HashMap::new())); @@ -208,44 +259,34 @@ pub async fn create_pty( cmd.env("LANG", "en_US.UTF-8"); cmd.env("LC_ALL", "en_US.UTF-8"); + // Enrolled below, and reaped by the guard until it reaches the registry. + #[allow(clippy::disallowed_methods)] let child = pair .slave .spawn_command(cmd) .map_err(|e| TerminalExtError::Internal(format!("failed to spawn shell: {e}")))?; - // Until the session reaches the registry nothing else can reach this shell, - // so every failure below has to kill it here or it is orphaned. - let mut shell = Shell::Running { child, group: None }; + let mut shell = UnregisteredShell::new(child); if let Some(pid) = shell.pid() { - match xai_tty_utils::global_process_scope().enroll_terminal_pid(pid) { - Ok(enrolled) => shell.attach_group(enrolled), - Err(e) => { - shell.reap_now(); - return Err(TerminalExtError::Internal(format!( - "failed to enroll shell: {e}" - ))); - } - } + // `enroll_terminal_pid` reaps with a grace wait if it loses the close + // race, so run it off-task. + let enrolled = tokio::task::spawn_blocking(move || { + xai_tty_utils::global_process_scope().enroll_terminal_pid(pid) + }) + .await + .map_err(|e| TerminalExtError::Internal(format!("enroll task failed: {e}")))? + .map_err(|e| TerminalExtError::Internal(format!("failed to enroll shell: {e}")))?; + shell.attach_group(enrolled); } - let reader = match pair.master.try_clone_reader() { - Ok(reader) => reader, - Err(e) => { - shell.reap_now(); - return Err(TerminalExtError::Internal(format!( - "failed to clone pty reader: {e}" - ))); - } - }; - let writer = match pair.master.take_writer() { - Ok(writer) => writer, - Err(e) => { - shell.reap_now(); - return Err(TerminalExtError::Internal(format!( - "failed to take pty writer: {e}" - ))); - } - }; + let reader = pair + .master + .try_clone_reader() + .map_err(|e| TerminalExtError::Internal(format!("failed to clone pty reader: {e}")))?; + let writer = pair + .master + .take_writer() + .map_err(|e| TerminalExtError::Internal(format!("failed to take pty writer: {e}")))?; let (input_tx, input_rx) = mpsc::channel(INPUT_CHANNEL_CAPACITY); spawn_pty_input_loop(writer, input_rx); @@ -277,12 +318,24 @@ pub async fn create_pty( }) }); - let mut session = PtySession { + // Taking the lock first means nothing can await between disarming the guard + // and the insert that gives teardown another way to reach the shell. + let mut registry = PTY_REGISTRY.lock().await; + + // The scope can close during the setup above, and teardown has already run + // by then: publishing here would advertise a shell it just killed. + if xai_tty_utils::global_process_scope().is_closed() { + return Err(TerminalExtError::Internal( + "process scope closed while the shell was starting".to_string(), + )); + } + + let entry = Arc::new(Mutex::new(PtySession { master: Some(pair.master), input_tx: Some(input_tx), output_offset: 0, output_ring: VecDeque::with_capacity(OUTPUT_RING_BUFFER_SIZE), - shell, + shell: shell.into_registered(), cwd: resolved_cwd, name: resolved_name, created_at, @@ -291,22 +344,9 @@ pub async fn create_pty( target_client_id, busy: false, gateway: gateway.clone(), - }; - - // The scope can close during the setup above, and teardown has already run - // by then: publishing here would advertise a shell it just killed. - if xai_tty_utils::global_process_scope().is_closed() { - session.shell.reap_now(); - return Err(TerminalExtError::Internal( - "process scope closed while the shell was starting".to_string(), - )); - } - - let entry = Arc::new(Mutex::new(session)); - PTY_REGISTRY - .lock() - .await - .insert(pty_id.clone(), entry.clone()); + })); + registry.insert(pty_id.clone(), entry.clone()); + drop(registry); let pty_id_clone = pty_id.clone(); tokio::task::spawn_local(run_pty_output_loop(reader, entry, pty_id_clone, gateway)); @@ -447,21 +487,11 @@ async fn run_pty_output_loop( ); } -/// Whether the PTY's controlling terminal has a foreground process group -/// distinct from the shell itself — i.e. a command is actively running -/// rather than the shell sitting idle at its prompt. +/// Whether a command is running, rather than the shell sitting at its prompt. /// -/// `process_group_leader()` issues `tcgetpgrp` on the master fd; an idle -/// shell is its own foreground process group, so it matches the shell -/// child's pid. When a command runs in the foreground the kernel reports -/// that command's process group instead. Returns false when the value is -/// unavailable (the shell exited or runs without job control). -/// -/// Limitation: a shell that `exec`s a program in place keeps the same pid and -/// pgid, so `tcgetpgrp` still matches the recorded child pid and the program -/// reads as idle. Telling that apart from a real idle prompt needs per-OS -/// process inspection, so a command launched the usual way (fork then exec) is -/// detected while an `exec`-replaced shell is not. +/// A shell that `exec`s a program in place keeps its pid and pgid, so that +/// program reads as idle. Separating it from a real prompt needs per-OS +/// process inspection. #[cfg(unix)] fn session_has_foreground_process(session: &PtySession) -> bool { let Some(foreground_pgid) = session @@ -642,12 +672,8 @@ pub async fn close_all() { } } -/// Resolve the shell binary and arguments for an interactive PTY session. -/// -/// Priority: explicit `shell` param > `$SHELL` env > platform default. -/// On Windows falls back to the `detect_windows_shell` cascade -/// (pwsh > powershell.exe > Git Bash > cmd.exe, overridable via -/// `GROK_SHELL`) since `$SHELL` is absent. +/// Explicit `shell` param, then `$SHELL`, then the platform default. Windows +/// has no `$SHELL`, so it uses the `detect_windows_shell` cascade. fn resolve_pty_shell(shell: Option<&str>) -> (String, Vec) { if let Some(s) = shell { return (s.to_string(), vec![]); @@ -712,12 +738,9 @@ pub struct PtyLoadResult { pub exit_code: Option, } -/// Reconnect to a PTY. Replays the full ring buffer (with `isReplay: true`) -/// so the client can reset its VTE emulator and feed all bytes from scratch. -/// Exited PTYs are still loadable so the client can see final output. -/// -/// Updates the stored `target_client_id` so that subsequent output -/// notifications from the output loop are routed to the reconnecting client. +/// Reconnect to a PTY, replaying the ring buffer so the client can reset its +/// VTE emulator and feed all bytes from scratch. An exited PTY still loads, so +/// its final output stays readable. pub async fn load( pty_id: &str, gateway: &GatewaySender, @@ -1005,6 +1028,36 @@ mod tests { .await; } + /// Covers the failure paths in [`create_pty`], which reap by returning. + #[tokio::test] + async fn dropping_an_unregistered_shell_reaps_it() { + let pair = native_pty_system() + .openpty(PtySize { + rows: 24, + cols: 80, + pixel_width: 0, + pixel_height: 0, + }) + .expect("openpty"); + let mut cmd = CommandBuilder::new("/bin/sh"); + cmd.arg("-c"); + cmd.arg("sleep 300"); + #[allow(clippy::disallowed_methods)] + let child = pair.slave.spawn_command(cmd).expect("spawn shell"); + let pid = child.process_id().expect("shell pid") as i32; + + drop(UnregisteredShell::new(child)); + + let deadline = std::time::Instant::now() + Duration::from_secs(5); + while unsafe { libc::kill(pid, 0) } == 0 { + assert!( + std::time::Instant::now() < deadline, + "shell {pid} survived the guard drop" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + #[cfg(unix)] async fn wait_for_reported_pid(pty_id: &str) -> i32 { let deadline = std::time::Instant::now() + Duration::from_secs(10); diff --git a/crates/codegen/xai-grok-shell/src/test_support/mod.rs b/crates/codegen/xai-grok-shell/src/test_support/mod.rs index d5b8863..1136dba 100644 --- a/crates/codegen/xai-grok-shell/src/test_support/mod.rs +++ b/crates/codegen/xai-grok-shell/src/test_support/mod.rs @@ -2,6 +2,15 @@ pub(crate) mod lsp_runtime; pub(crate) const TEST_MODEL: &str = "test-model"; +/// Keep this crate's unit-test binary from writing synthetic events into +/// the real unified log; pre-main so the redirect beats the lazily-opened +/// writer. Integration binaries under `tests/` isolate via `TestSandbox` +/// homes instead. +#[ctor::ctor] +fn redirect_unified_log_for_tests() { + xai_grok_telemetry::unified_log::redirect_to_temp_for_tests(); +} + /// Prepend the hermetic git binary (via `GIT_BIN_PATH`) to `PATH` so that /// `Command::new("git")` in test helpers resolves to the Bazel-provided /// static binary instead of relying on system-installed git. diff --git a/crates/codegen/xai-grok-shell/src/util/dual_clock.rs b/crates/codegen/xai-grok-shell/src/util/dual_clock.rs new file mode 100644 index 0000000..b50ece4 --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/util/dual_clock.rs @@ -0,0 +1,52 @@ +//! One instant captured on two clocks, so elapsed time stays honest across +//! a system suspend without trusting the wall clock alone. + +use std::time::{Duration, Instant, SystemTime}; + +/// `Instant` is monotonic but, on macOS (`mach_absolute_time`) and Linux +/// (`CLOCK_MONOTONIC`), *pauses while the machine is asleep*, so it alone +/// under-reports any span containing a suspend. `SystemTime` keeps advancing +/// through sleep but jumps with NTP steps and manual changes. Capturing both +/// lets a caller bound elapsed *awake* time (mono), elapsed *real* time +/// (wall), and their difference — which grows by exactly the suspended time. +#[derive(Clone, Copy)] +pub(crate) struct DualClock { + /// Monotonic; pauses during sleep. Bounds elapsed *awake* time. + pub(crate) mono: Instant, + /// Wall clock; advances through sleep. Bounds elapsed *real* time. + pub(crate) wall: SystemTime, +} + +impl DualClock { + pub(crate) fn now() -> Self { + Self { + mono: Instant::now(), + wall: SystemTime::now(), + } + } + + /// Elapsed on each clock as `(monotonic, wall)`. Wall elapsed clamps to + /// zero if the clock ran backwards (NTP step) so a backward jump can + /// never fabricate a suspend or inflate a duration. + pub(crate) fn elapsed_between(&self, now: DualClock) -> (Duration, Duration) { + ( + now.mono.saturating_duration_since(self.mono), + now.wall.duration_since(self.wall).unwrap_or(Duration::ZERO), + ) + } + + /// [`Self::elapsed_between`] against the live clocks. + pub(crate) fn elapsed(&self) -> (Duration, Duration) { + self.elapsed_between(Self::now()) + } + + /// `(awake, total, suspended)` durations since this instant. + pub(crate) fn elapsed_split(&self) -> (Duration, Duration, Duration) { + let (awake, total) = self.elapsed(); + (awake, total, total.saturating_sub(awake)) + } +} + +#[cfg(test)] +#[path = "dual_clock_tests.rs"] +mod tests; diff --git a/crates/codegen/xai-grok-shell/src/util/dual_clock_tests.rs b/crates/codegen/xai-grok-shell/src/util/dual_clock_tests.rs new file mode 100644 index 0000000..c25a7b0 --- /dev/null +++ b/crates/codegen/xai-grok-shell/src/util/dual_clock_tests.rs @@ -0,0 +1,18 @@ +use std::time::Duration; + +use super::DualClock; + +/// A backward wall jump (NTP step) clamps to zero rather than underflowing, +/// so it can never fabricate a suspend or inflate a duration. +#[test] +fn backward_wall_jump_clamps_to_zero() { + let start = DualClock::now(); + let stepped_back = DualClock { + mono: start.mono + Duration::from_secs(5), + wall: start.wall - Duration::from_secs(60), + }; + assert_eq!( + start.elapsed_between(stepped_back), + (Duration::from_secs(5), Duration::ZERO) + ); +} diff --git a/crates/codegen/xai-grok-shell/src/util/mod.rs b/crates/codegen/xai-grok-shell/src/util/mod.rs index de0f0cf..69ed5e5 100644 --- a/crates/codegen/xai-grok-shell/src/util/mod.rs +++ b/crates/codegen/xai-grok-shell/src/util/mod.rs @@ -1,4 +1,5 @@ pub mod config; +pub(crate) mod dual_clock; pub mod grok_auth_credentials; pub mod hooks; pub mod limits; diff --git a/crates/codegen/xai-grok-shell/tests/test_leader_soak.rs b/crates/codegen/xai-grok-shell/tests/test_leader_soak.rs index 49513b9..dc6712e 100644 --- a/crates/codegen/xai-grok-shell/tests/test_leader_soak.rs +++ b/crates/codegen/xai-grok-shell/tests/test_leader_soak.rs @@ -105,6 +105,8 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() { let server = xai_grok_test_support::MockInferenceServer::start() .await .unwrap(); + // Measure the leader, not the harness's copy of every conversation. + server.set_keep_requests(false); let grok_home = TempDir::new().unwrap(); let workdir = TempDir::new().unwrap(); @@ -207,8 +209,6 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() { let mut turns: u64 = 0; let mut baseline: Option = None; - // Each cycle: 10 fresh clients, 2 sessions each, one scripted - // turn per session, then all disconnect. while tokio::time::Instant::now() < soak_deadline { cycles += 1; let mut clients = Vec::new(); @@ -282,7 +282,7 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() { } // An entry that never drains names itself here, one cycle - // after it leaks, while memory is still within its budget. + // after it leaks. let counts = registry_counts(&mut bootstrap, 1000 + cycles).await; assert_eq!( counts["sessions"], 0, @@ -332,7 +332,7 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() { ({:.2} MB per cycle)", net_bytes as f64 / measured as f64 / (1024.0 * 1024.0) ); - let max_per_cycle = env_u64("LEADER_SOAK_MAX_HEAP_BYTES_PER_CYCLE", 4 << 20) as i64; + let max_per_cycle = env_u64("LEADER_SOAK_MAX_HEAP_BYTES_PER_CYCLE", 1 << 20) as i64; assert!( per_cycle <= max_per_cycle, "leader retained {per_cycle} heap bytes per cycle (bound {max_per_cycle})" diff --git a/crates/codegen/xai-grok-shell/tests/test_registry_churn.rs b/crates/codegen/xai-grok-shell/tests/test_registry_churn.rs index 21e684e..624ead8 100644 --- a/crates/codegen/xai-grok-shell/tests/test_registry_churn.rs +++ b/crates/codegen/xai-grok-shell/tests/test_registry_churn.rs @@ -31,6 +31,7 @@ const RPC_TIMEOUT: Duration = Duration::from_secs(60); #[serde(deny_unknown_fields)] struct Counts { sessions: usize, + loading_sessions: usize, session_threads: usize, resident_resources: usize, retained_resources: usize, @@ -45,6 +46,7 @@ struct Counts { subagent_active: usize, subagent_completed: usize, workspace_bindings: Option, + workspace_activity_sessions: Option, } struct AutoApproveClient; #[async_trait::async_trait(?Send)] @@ -91,6 +93,19 @@ async fn read_counts(conn: &acp::ClientSideConnection) -> Counts { serde_json::from_value(resp["result"]["registries"].clone()) .unwrap_or_else(|e| panic!("x.ai/debug/agent: bad registries payload: {e}\n{resp}")) } +/// Counts read once the actor threads are reaped. Nothing signals a thread +/// exit, so this polls; both ends settle, so neither catches one mid-exit. +async fn settled_counts(conn: &acp::ClientSideConnection) -> Counts { + let mut counts = read_counts(conn).await; + for _ in 0..100 { + if counts.session_threads == 0 { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + counts = read_counts(conn).await; + } + counts +} async fn new_session(conn: &acp::ClientSideConnection, cwd: &std::path::Path) -> acp::SessionId { tokio::time::timeout( RPC_TIMEOUT, @@ -252,21 +267,29 @@ fn session_churn_returns_registry_snapshot_to_baseline() { agent_rt.block_on(local.run_until(async move { let client_conn = connect_and_auth().await; churn_one(&client_conn, workdir.path(), 0).await; - let baseline = read_counts(&client_conn).await; + let baseline = settled_counts(&client_conn).await; assert_eq!( baseline.sessions, 0, "warmup session must be fully removed before baseline" ); assert_eq!( - (baseline.resident_resources, baseline.retained_resources), - (0, 0), + ( + baseline.resident_resources, + baseline.retained_resources, + baseline.loading_sessions + ), + (0, 0, 0), "warmup must leave no per-session resource entries, including \ entries holding no resources" ); assert_eq!( - baseline.workspace_bindings, - Some(0), - "warmup must have built the local workspace and released its binding" + ( + baseline.workspace_bindings, + baseline.workspace_activity_sessions + ), + (Some(0), Some(0)), + "warmup must have built the local workspace and released both its \ + binding and its activity record" ); assert_eq!( ( @@ -295,7 +318,7 @@ fn session_churn_returns_registry_snapshot_to_baseline() { })) .await; futures::future::join_all(concurrent.iter().map(|sid| close_session(conn, sid))).await; - let after = read_counts(&client_conn).await; + let after = settled_counts(&client_conn).await; assert_eq!( after, baseline, "session churn must return every registry count to baseline \ diff --git a/crates/codegen/xai-grok-shell/tests/test_sampling_client.rs b/crates/codegen/xai-grok-shell/tests/test_sampling_client.rs index 45db3ea..470cf3d 100644 --- a/crates/codegen/xai-grok-shell/tests/test_sampling_client.rs +++ b/crates/codegen/xai-grok-shell/tests/test_sampling_client.rs @@ -896,7 +896,7 @@ async fn test_chat_completions_401_unauthorized() { let result = client.conversation_stream(request).await; assert!(result.is_err()); - if let Err(SamplingError::Auth(_)) = result { + if let Err(SamplingError::Auth { .. }) = result { // Expected } else { panic!("Expected Auth error"); @@ -939,7 +939,7 @@ async fn test_responses_api_401_unauthorized() { let result = client.conversation_stream_responses(request).await; assert!(result.is_err()); - if let Err(SamplingError::Auth(_)) = result { + if let Err(SamplingError::Auth { .. }) = result { // Expected } else { panic!("Expected Auth error"); diff --git a/crates/codegen/xai-grok-telemetry/Cargo.toml b/crates/codegen/xai-grok-telemetry/Cargo.toml index bf42f19..3eab3e9 100644 --- a/crates/codegen/xai-grok-telemetry/Cargo.toml +++ b/crates/codegen/xai-grok-telemetry/Cargo.toml @@ -80,6 +80,8 @@ tokio-stream = { workspace = true } tracing-opentelemetry = { workspace = true } [dev-dependencies] +# Pre-main unified-log redirect for this crate's own test binary. +ctor = { workspace = true } tonic = { workspace = true, features = ["transport"] } # In-memory log/metric exporters for the external-stream wire-shape tests. opentelemetry_sdk = { workspace = true, features = ["testing"] } diff --git a/crates/codegen/xai-grok-telemetry/src/events.rs b/crates/codegen/xai-grok-telemetry/src/events.rs index 6454d4e..4ca5f1e 100644 --- a/crates/codegen/xai-grok-telemetry/src/events.rs +++ b/crates/codegen/xai-grok-telemetry/src/events.rs @@ -289,14 +289,30 @@ pub struct LoginCompleted { pub mid_session: bool, } -/// A login flow failed. `error` is the raw error message from the auth flow. +/// How a login attempt's HTTP request failed. +#[derive(Debug, Serialize, Clone, Copy, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum LoginFailureKind { + /// `is_connect`: a dead TCP connect *or* a TLS handshake killed + /// mid-flight. `os_error` tells them apart. + TransportConnect, + /// In-flight request cut short: reset, close, timeout, body phase. + TransportInterrupted, + /// Client-side request construction / redirect policy defect. + TransportPermanent, + Decode, +} + +/// One per failed login attempt, emitted by the login funnel so a retried +/// request can't inflate the count. Failures that never reached HTTP (user +/// backed out, loopback bind, id_token validation) are not reported. #[derive(Serialize)] pub struct LoginFailed { - pub method: String, - pub mode: String, + pub error_kind: LoginFailureKind, + /// OS code from the failure's cause chain (54/104 ECONNRESET, 10054 on + /// Windows). #[serde(skip_serializing_if = "Option::is_none")] - pub error: Option, - pub duration_ms: u64, + pub os_error: Option, } /// The user backed out of the login funnel. `stage` is "picker", @@ -2204,6 +2220,29 @@ mod tests { ); } + #[test] + fn login_failed_serializes_kind_and_os_code() { + let v = serde_json::to_value(LoginFailed { + error_kind: LoginFailureKind::TransportInterrupted, + os_error: Some(104), + }) + .unwrap(); + assert_eq!( + v, + serde_json::json!({ "error_kind": "transport_interrupted", "os_error": 104 }) + ); + } + + #[test] + fn login_failed_omits_absent_os_code() { + let v = serde_json::to_value(LoginFailed { + error_kind: LoginFailureKind::Decode, + os_error: None, + }) + .unwrap(); + assert_eq!(v, serde_json::json!({ "error_kind": "decode" })); + } + #[test] fn api_key_save_result_omits_error_when_ok() { let ok = serde_json::to_value(ApiKeySaveResult { diff --git a/crates/codegen/xai-grok-telemetry/src/session_ctx.rs b/crates/codegen/xai-grok-telemetry/src/session_ctx.rs index e1659c3..666a911 100644 --- a/crates/codegen/xai-grok-telemetry/src/session_ctx.rs +++ b/crates/codegen/xai-grok-telemetry/src/session_ctx.rs @@ -5,6 +5,7 @@ //! Extracted from `xai-grok-shell::agent::telemetry`. use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use serde::Serialize; use serde_json::json; @@ -198,6 +199,37 @@ pub fn emit_event(event_suffix: impl Into emit_event_with_origin(EmitterOrigin::Shell, event_suffix, data); } +/// Posts spawned by [`emit_event_with_origin`] that haven't finished. Emission +/// is fire-and-forget so it never blocks a turn, which also means a process +/// exiting right after emitting drops the event — see [`drain_pending`]. +static PENDING_EVENTS: AtomicUsize = AtomicUsize::new(0); + +/// Decrement on every exit path, including a panicking or cancelled post. +struct PendingEventGuard; + +impl Drop for PendingEventGuard { + fn drop(&mut self) { + PENDING_EVENTS.fetch_sub(1, Ordering::Release); + } +} + +/// Wait (up to `timeout`) for in-flight event posts to finish. For commands +/// that exit as soon as their work is done; the agent runs long enough that +/// its events land on their own. +pub async fn drain_pending(timeout: std::time::Duration) { + let deadline = std::time::Instant::now() + timeout; + while PENDING_EVENTS.load(Ordering::Acquire) > 0 { + if std::time::Instant::now() >= deadline { + tracing::debug!( + pending = PENDING_EVENTS.load(Ordering::Acquire), + "telemetry: gave up draining pending events" + ); + return; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } +} + /// Emit an event whose analytics name is `{origin prefix}{event_suffix}`. pub fn emit_event_with_origin( origin: EmitterOrigin, @@ -214,7 +246,15 @@ pub fn emit_event_with_origin( }) .ok(); + if tokio::runtime::Handle::try_current().is_err() { + // `spawn` below panics without a runtime; counting first would pin the + // gauge above zero for the rest of the process. + tracing::debug!(event = %event_name, "telemetry: no runtime, dropping event"); + return; + } + PENDING_EVENTS.fetch_add(1, Ordering::Release); tokio::spawn(async move { + let _pending = PendingEventGuard; let user_ctx = UserContext::collect(); let request_id = format!("{}-{}", event_name, uuid::Uuid::new_v4()); @@ -264,6 +304,30 @@ mod tests { }); } + /// What a command exiting right after emitting (`grok login`) relies on. + /// Asserts on the wait, not on the gauge: it is process-global and other + /// tests in this binary emit concurrently. + #[tokio::test] + async fn drain_pending_waits_for_in_flight_posts() { + emit_event_with_origin( + EmitterOrigin::Shell, + "drain_probe", + json!({ "probe": true }), + ); + assert!( + PENDING_EVENTS.load(Ordering::Acquire) > 0, + "emission must register before the post is awaited" + ); + + let started = std::time::Instant::now(); + let budget = std::time::Duration::from_secs(5); + drain_pending(budget).await; + assert!( + started.elapsed() < budget, + "drain must observe the post finish, not time out" + ); + } + /// Event-name prefixes are wire contract — analytics queries match on them, so /// they must not drift. #[test] diff --git a/crates/codegen/xai-grok-telemetry/src/unified_log.rs b/crates/codegen/xai-grok-telemetry/src/unified_log.rs index 4c27ab9..bf2a719 100644 --- a/crates/codegen/xai-grok-telemetry/src/unified_log.rs +++ b/crates/codegen/xai-grok-telemetry/src/unified_log.rs @@ -167,10 +167,69 @@ type FileIdentity = (u64, u64); static WRITER: LazyLock>> = LazyLock::new(|| Mutex::new(open_writer())); +/// See [`redirect_to_temp_for_tests`]. +static TEST_REDIRECT: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); + +/// Redirect all subsequent unified-log writes **and** snapshot reads to a +/// per-process file under the system temp directory, so test binaries stop +/// writing synthetic events into the developer's real +/// `~/.grok/logs/unified.jsonl` (those bursts inflate exactly the counters +/// an incident responder greps for). Runtime-activated rather than a cargo +/// feature: Bazel compiles production and test targets with one shared +/// feature set, so a feature gate would leak into production builds. +/// +/// Idempotent and safe at any point: an already-open writer is re-pointed, +/// so an emit that precedes the redirect cannot pin the real path. Test +/// binaries install it pre-main via `#[ctor]`. +pub fn redirect_to_temp_for_tests() { + TEST_REDIRECT.store(true, std::sync::atomic::Ordering::Relaxed); + if let Ok(mut guard) = WRITER.lock() { + *guard = open_writer(); + } +} + fn log_path() -> PathBuf { + if TEST_REDIRECT.load(std::sync::atomic::Ordering::Relaxed) { + return test_log_dir().join(LOG_FILE); + } grok_home().join(LOG_DIR).join(LOG_FILE) } +/// Owner-only (0o700), freshly-created directory for the test redirect. +/// +/// The stream carries path metadata and credential tail fragments, and the +/// system temp dir is world-writable on Linux: a pre-planted directory or +/// symlink would let another local user read the file — or make the writer +/// and [`trim_file`] operate through a symlink onto a victim file. The +/// non-recursive `create` fails on any pre-existing path instead of +/// adopting it, and the nanos component makes the name unpredictable. +/// Panicking on failure is deliberate: this branch only runs in test +/// binaries, and silently falling back would reopen the hole via +/// `open_writer_at`'s `create_dir_all`. +fn test_log_dir() -> &'static PathBuf { + static TEST_LOG_DIR: OnceLock = OnceLock::new(); + TEST_LOG_DIR.get_or_init(|| { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0); + let dir = std::env::temp_dir().join(format!( + "grok-unified-log-test-{}-{nanos}", + std::process::id() + )); + let mut builder = fs::DirBuilder::new(); + #[cfg(unix)] + { + use std::os::unix::fs::DirBuilderExt; + builder.mode(0o700); + } + builder + .create(&dir) + .expect("create private unified-log test dir"); + dir + }) +} + pub fn file_size(path: &std::path::Path) -> u64 { fs::metadata(path).map(|m| m.len()).unwrap_or(0) } @@ -516,6 +575,34 @@ pub fn snapshot_session_log(session_id: &str) -> Option> { mod tests { use super::*; + /// Pre-main, so no test in this binary can race the lazily-opened + /// writer onto the developer's real `~/.grok/logs/unified.jsonl`. + #[ctor::ctor] + fn redirect_for_tests() { + redirect_to_temp_for_tests(); + } + + /// The redirect must cover both the writer and the snapshot readers: + /// an emit lands in a per-process temp file, never under `grok_home()`. + #[test] + fn redirect_routes_writes_and_snapshots_to_process_temp_file() { + info( + "unified-log redirect probe", + Some("redirect-probe-sid"), + None, + ); + let snapshot = snapshot_log().expect("snapshot after emit"); + assert!( + String::from_utf8_lossy(&snapshot).contains("unified-log redirect probe"), + "snapshot must read the same redirected file the writer appended to" + ); + assert!( + log_path().starts_with(std::env::temp_dir()), + "the shared file must live under the temp dir, not grok_home(): {}", + log_path().display() + ); + } + #[test] fn log_entry_serializes_minimal() { let entry = LogEntry { diff --git a/crates/codegen/xai-grok-test-support/src/mock_server.rs b/crates/codegen/xai-grok-test-support/src/mock_server.rs index fe50a16..10eaaa1 100644 --- a/crates/codegen/xai-grok-test-support/src/mock_server.rs +++ b/crates/codegen/xai-grok-test-support/src/mock_server.rs @@ -1,13 +1,13 @@ -//! Mock inference server with request logging and automatic cleanup. +//! Mock inference server. Logs every request and shuts down on drop. //! -//! Serves `/v1/chat/completions`, `/v1/responses`, and `/v1/messages` in one -//! of two response modes: echo (default — streams `Echo: `) -//! or a fixed text set via [`MockInferenceServer::set_response`] (streamed -//! with byte-exact reconstruction). Named request-matched expectations take -//! precedence, followed by compatibility per-path [`ScriptedResponse`] FIFOs. -//! `/v1/models` and `/v1/settings` return -//! configurable responses (settings is 404 until set). All requests are -//! logged — bodies and headers — for assertion in tests. +//! Serves the three inference endpoints (`/v1/chat/completions`, +//! `/v1/responses`, `/v1/messages`) plus `/v1/models`, `/v1/settings`, +//! `/v1/user`, `/v1/storage`, and `/v1/privacy/coding-data-retention`. +//! +//! The inference endpoints answer from the first source that matches: a named +//! expectation, then the path's [`ScriptedResponse`] queue, then the active +//! mode, which echoes the last user message until +//! [`MockInferenceServer::set_response`] replaces it with a fixed text. use std::collections::VecDeque; use std::convert::Infallible; @@ -40,10 +40,8 @@ pub struct LogEntry { pub method: String, pub path: String, pub body: Option, - /// Value of the `Authorization` header, if present. pub authorization: Option, - /// Request headers (lowercase names, arrival order), captured on the - /// inference POST endpoints; the GET endpoints log an empty list. + /// Lowercase names in arrival order. Empty for the GET endpoints. pub headers: Vec<(String, String)>, } @@ -58,13 +56,13 @@ impl LogEntry { } } -/// Requests kept for inspection; each holds a whole conversation, so an -/// unbounded log outgrows what it is testing. `request_count` stays exact. +/// An entry holds a whole conversation, so the log evicts oldest first. const MAX_LOGGED_REQUESTS: usize = 1024; pub struct RequestLog { count: AtomicU32, entries: std::sync::Mutex>, + keep_entries: AtomicBool, } impl RequestLog { @@ -72,6 +70,7 @@ impl RequestLog { Self { count: AtomicU32::new(0), entries: std::sync::Mutex::new(Vec::new()), + keep_entries: AtomicBool::new(true), } } @@ -84,6 +83,9 @@ impl RequestLog { headers: Vec<(String, String)>, ) { self.count.fetch_add(1, Ordering::SeqCst); + if !self.keep_entries.load(Ordering::SeqCst) { + return; + } let mut entries = self.entries.lock().unwrap(); if entries.len() >= MAX_LOGGED_REQUESTS { entries.remove(0); @@ -98,26 +100,19 @@ impl RequestLog { } } -/// A model entry for the mock `/v1/models` endpoint. +/// A model served by `/v1/models`. Each field is emitted under its camelCase +/// name when set, at the top level except for `agent_type`, which goes in +/// `_meta`. #[derive(Debug, Clone)] pub struct MockModelEntry { - /// Model ID (e.g. `"test-model"`). pub id: String, - /// Optional agent type (e.g. `"cursor"`). - /// Emitted as `agentType` inside `_meta` when set. pub agent_type: Option, - /// Optional API backend (e.g. `"messages"`). Emitted as `apiBackend` - /// when set; absent means the shell's default backend. pub api_backend: Option, - /// Emitted as `supportsBackendSearch` when true. pub supports_backend_search: bool, - /// Emitted as `supportsReasoningEffort` (top-level) when true. pub supports_reasoning_effort: bool, - /// Emitted as `reasoningEffort` (top-level) when set. pub reasoning_effort: Option, - /// Emitted as `reasoningEfforts` (top-level) when non-empty. Each entry is a - /// raw JSON option (a table `{ "value": ..., "id"?, "label"?, ... }` or a - /// bare value string), matching what `parse_remote_model_value` reads. + /// Each entry is a table carrying a `value` key, or a bare value string. + /// `parse_remote_model_value` defines the full shape. pub reasoning_efforts: Vec, } @@ -195,12 +190,10 @@ impl MockModelEntry { } } -/// What the inference endpoints stream back. enum ResponseMode { - /// Echo the last user message as `Echo: ` (whitespace-collapsing). + /// `Echo: `, with whitespace collapsed. Echo, - /// Stream a fixed text whose deltas reconstruct it byte-for-byte - /// (newlines preserved — required for fenced code blocks). + /// Deltas reconstruct the text byte for byte, newlines included. Fixed(String), } @@ -228,24 +221,19 @@ fn paced_events( ) } -/// Max body bytes retained on each accepted [`StorageUpload`] (keeps large -/// e2e artifacts from ballooning test memory; meta/small dumps stay intact). const STORAGE_BODY_CAPTURE_CAP: usize = 256 * 1024; -/// One accepted (HTTP 200) mock `/v1/storage` upload. +/// An upload `/v1/storage` accepted. #[derive(Debug, Clone)] pub struct StorageUpload { pub path: String, pub size: usize, - /// Request body when `size <= 256 KiB`; empty for larger payloads. + /// Empty when `size` exceeds `STORAGE_BODY_CAPTURE_CAP`. pub body: Vec, - /// `Authorization` header value as sent (e.g. `Bearer …`). pub authorization: Option, } -/// Mock `/v1/storage` state: a flippable 401 gate plus a record of accepted -/// uploads, so e2e tests can simulate an auth outage window and assert the -/// trace upload queue parks, then drains after the gate heals. +/// A 401 gate tests can flip, plus the uploads accepted through it. #[derive(Default)] struct StorageState { unauthorized: AtomicBool, @@ -253,10 +241,6 @@ struct StorageState { uploads: std::sync::Mutex>, } -/// Mock `/v1/chat/completions` + `/v1/responses` + `/v1/messages` + -/// `/v1/models` + `/v1/settings` + `/v1/storage` + -/// `/v1/privacy/coding-data-retention` server. -/// Logs all requests. Shuts down on drop. pub struct MockInferenceServer { addr: SocketAddr, shutdown_tx: Option>, @@ -265,31 +249,23 @@ pub struct MockInferenceServer { settings: Arc>>, response_mode: Arc>, overrides: InferenceOverrides, - /// Per-agent-turn assistant texts (see [`set_agent_turns`]). - /// - /// [`set_agent_turns`]: Self::set_agent_turns + /// One assistant text per agent turn, consumed in order. agent_turns: Arc>>, - /// `stop_reason` emitted by the `/v1/messages` terminal `message_delta`. + /// `stop_reason` on the `/v1/messages` terminal `message_delta`. messages_stop_reason: Arc>, - /// Optional per-SSE-event delay on all inference endpoints. chunk_delay: Arc>>, - /// Mock `/v1/storage` 401 gate + accepted-upload record. storage: Arc, - /// When set, `/v1/models` and `/v1/settings` hang forever (never - /// respond); see [`Self::set_hang`]. + /// When set, `/v1/models` and `/v1/settings` never respond. hang: Arc, - /// See [`Self::set_user_subscription_tier`]. user_tier: Arc>>, } impl MockInferenceServer { - /// Start with a single default `test-model` (no agent_type). + /// Serves one `test-model` with no agent type. pub async fn start() -> anyhow::Result { Self::start_with_models(vec![MockModelEntry::new("test-model")]).await } - /// Start with custom models. Use [`MockModelEntry::with_agent_type`] to - /// configure models with specific harness types for agent-type tests. pub async fn start_with_models(models: Vec) -> anyhow::Result { Self::start_inner(models, None).await } @@ -374,24 +350,18 @@ impl MockInferenceServer { }) } - /// Replace the model list at runtime. The next `/v1/models` request - /// (e.g. during session resume) will return the new list. pub fn set_models(&self, models: Vec) { let mut guard = self.models.write().unwrap(); *guard = models.iter().map(MockModelEntry::to_json).collect(); } - /// Stream this fixed text from all inference endpoints instead of echoing - /// the user message. Deltas reconstruct the text byte-for-byte (newlines - /// preserved). Subsequent calls replace the text. + /// Stream this text instead of echoing. Deltas reconstruct it byte for byte. pub fn set_response(&self, text: impl Into) { *self.response_mode.write().unwrap() = ResponseMode::Fixed(text.into()); } - /// Queue a [`ScriptedResponse`] for the next request on `path` (e.g. - /// `"/v1/chat/completions"`). Scripts are consumed FIFO per path by the - /// three inference endpoints; when a path's queue is empty, requests fall - /// back to the active response mode (echo/fixed). + /// Consumed FIFO per `path`, e.g. `"/v1/chat/completions"`. An empty queue + /// falls back to the response mode. pub fn enqueue_response(&self, path: impl Into, response: ScriptedResponse) { self.overrides.enqueue_response(path, response); } @@ -420,65 +390,54 @@ impl MockInferenceServer { .register_expectation(name, matcher, response, true) } - /// Queue one byte-exact response per foreground turn as compatibility sugar. + /// Queue one byte-exact response per foreground turn. pub fn set_agent_turns(&self, turns: impl IntoIterator) { *self.agent_turns.lock().unwrap() = turns.into_iter().collect(); } - /// Replace the settings at runtime. The next `GET /v1/settings` request - /// will return the new value as JSON. Until set, `/v1/settings` returns 404. + /// Until this is called, `GET /v1/settings` returns 404. pub fn set_settings(&self, settings: impl serde::Serialize) { let value = serde_json::to_value(settings).expect("serialize settings"); let mut guard = self.settings.write().unwrap(); *guard = Some(value); } - /// Preset `/v1/settings` to the minimal `{"allow_access": true}` payload - /// that opens the subscription gate (clients treat a missing field as - /// `false` and would sit on the upsell screen). + /// The smallest settings payload that opens the subscription gate. Without + /// it a client sits on the upsell screen. pub fn preset_allow_access(&self) { self.set_settings(json!({ "allow_access": true })); } - /// Make `/v1/models` and `/v1/settings` hang forever, standing in for a - /// black-holed backend in non-blocking-startup tests. + /// Stand in for a black-holed backend. pub fn set_hang(&self, hang: bool) { self.hang.store(hang, std::sync::atomic::Ordering::Release); } - /// Set the `subscriptionTier` served by `GET /v1/user`. `None` - /// (default) omits the field, which the shell treats as "no qualifying - /// subscription" (free tier). + /// The `subscriptionTier` on `GET /v1/user`. `None`, the default, omits the + /// field, which the shell reads as the free tier. pub fn set_user_subscription_tier(&self, tier: Option<&str>) { *self.user_tier.write().unwrap() = tier.map(str::to_owned); } - /// Set the `stop_reason` emitted by the `/v1/messages` terminal - /// `message_delta` (default `"end_turn"`). + /// Defaults to `"end_turn"`. pub fn set_messages_stop_reason(&self, stop_reason: impl Into) { *self.messages_stop_reason.write().unwrap() = stop_reason.into(); } - /// Pace all inference SSE streams: each event is emitted after `delay`. - /// `None` (default) restores instant streaming. Lets PTY e2e tests hold a - /// turn visibly "streaming" long enough to interact with it mid-flight - /// (e.g. Esc-cancel). Applies to requests started after the call. + /// Emit each SSE event after `delay`, so a test can hold a turn visibly + /// streaming. `None` restores instant streaming. Applies to requests + /// started after the call. pub fn set_chunk_delay(&self, delay: Option) { *self.chunk_delay.write().unwrap() = delay; } - /// Hold foreground terminal SSE events until [`release_agent_completions`]. - /// Compatibility API; per-expectation blocking gives tighter ownership. - /// - /// [`release_agent_completions`]: Self::release_agent_completions + /// Hold foreground terminal SSE events until + /// [`Self::release_agent_completions`]. Prefer per-expectation blocking. pub fn hold_agent_completions(&self) { self.overrides.hold_completions(); } - /// Release a hold set by [`hold_agent_completions`], letting held (and - /// future) agent turns emit their terminal event and complete. - /// - /// [`hold_agent_completions`]: Self::hold_agent_completions + /// Let held and future agent turns emit their terminal event. pub fn release_agent_completions(&self) { self.overrides.release_completions(); } @@ -492,6 +451,11 @@ impl MockInferenceServer { self.log.count.load(Ordering::SeqCst) } + /// Stop retaining entries. [`Self::request_count`] stays exact. + pub fn set_keep_requests(&self, enabled: bool) { + self.log.keep_entries.store(enabled, Ordering::SeqCst); + } + pub fn requests(&self) -> Vec { self.log.entries.lock().unwrap().clone() } @@ -576,8 +540,7 @@ impl MockInferenceServer { }) } - /// Flip the mock `/v1/storage` 401 gate. While `true`, every upload is - /// rejected with 401 (the auth-outage window the park-on-401 e2e drives). + /// While closed, every `/v1/storage` upload is rejected with 401. pub fn set_storage_unauthorized(&self, unauthorized: bool) { self.storage .unauthorized @@ -589,14 +552,13 @@ impl MockInferenceServer { self.storage.request_count.load(Ordering::SeqCst) } - /// Snapshot of accepted (HTTP 200) `/v1/storage` uploads. + /// Only the uploads that were accepted. pub fn storage_uploads(&self) -> Vec { self.storage.uploads.lock().unwrap().clone() } - /// Mock `/v1/storage` upload: count the attempt, reject with 401 while the - /// gate is closed, else record the upload and mirror the proxy's - /// `UploadResponse` JSON shape. + /// Counts the attempt, then either rejects it or records it and answers in + /// the proxy's `UploadResponse` shape. fn storage_upload_handler( storage: &StorageState, headers: &HeaderMap, @@ -974,10 +936,8 @@ impl MockInferenceServer { if hang.load(std::sync::atomic::Ordering::Acquire) { tokio::time::sleep(Duration::from_secs(3600)).await; } - // Scripted one-shots take precedence (FIFO), so a - // test can serve a transient payload (e.g. one - // stale gated snapshot) and fall back to the - // steady-state `set_settings` value afterwards. + // Scripts take precedence, so a test can serve a + // transient payload before the steady-state value. if let Some(s) = overrides.pop_scripted("/v1/settings") { return s.into_response_paced(None, None).await; } @@ -1023,9 +983,8 @@ impl MockInferenceServer { let log = log.clone(); let user_tier = user_tier.clone(); async move { - // Keep the query string in the log so tests can - // count `?include=subscription` checks separately - // from plain enrichment fetches. + // Log the query string so a test can count + // `?include=subscription` on its own. let path = match query { Some(q) if !q.is_empty() => format!("/v1/user?{q}"), _ => "/v1/user".to_owned(), @@ -1054,9 +1013,8 @@ impl MockInferenceServer { } }), ) - // The shell probes these before/alongside per-file uploads. Answer - // 404 ("old proxy") so it falls back to plain `POST /v1/storage`, - // which is the path the park-on-401 e2e exercises. + // 404 reads as an old proxy, so the shell falls back to a plain + // `POST /v1/storage`. .route( "/v1/storage/exists", get(|| async { StatusCode::NOT_FOUND }), @@ -1266,6 +1224,25 @@ mod tests { ); } + #[tokio::test] + async fn dropping_entries_keeps_the_count_exact() { + let server = MockInferenceServer::start().await.unwrap(); + post_chat(&server, "kept").await; + let counted = server.request_count(); + let kept = server.requests().len(); + + server.set_keep_requests(false); + post_chat(&server, "dropped").await; + + assert_eq!(server.request_count(), counted + 1); + let entries = server.requests(); + assert_eq!(entries.len(), kept); + assert!( + format!("{:?}", entries.last().unwrap().body).contains("kept"), + "the surviving entry should be the one recorded before the switch" + ); + } + #[tokio::test] async fn auxiliary_request_does_not_consume_foreground_expectation() { let server = MockInferenceServer::start().await.unwrap(); diff --git a/crates/codegen/xai-grok-tools/Cargo.toml b/crates/codegen/xai-grok-tools/Cargo.toml index 813bf8d..6b447af 100644 --- a/crates/codegen/xai-grok-tools/Cargo.toml +++ b/crates/codegen/xai-grok-tools/Cargo.toml @@ -11,6 +11,7 @@ arc-swap = { workspace = true } dirs = { workspace = true } derive_more = { workspace = true, features = ["from", "try_into"] } dunce = { workspace = true } +encoding_rs = { workspace = true, optional = true } fs2 = { workspace = true } educe = { workspace = true, features = ["Debug"] } async-openai = { workspace = true } @@ -31,6 +32,7 @@ ignore = { workspace = true } infer = { workspace = true } globset = { workspace = true } image = { workspace = true, features = ["png", "jpeg", "gif", "webp", "bmp", "tiff", "ico"] } +kamadak-exif = { workspace = true, optional = true } pdf_oxide = { workspace = true } # PPTX text extraction (implementations/read_file/pptx.rs). Deliberately NOT # `workspace = true`: the workspace `zip` pin has default features on, which @@ -73,6 +75,7 @@ tokio = { workspace = true, features = [ tokio-util = { workspace = true, features = ["compat"] } tonic = { workspace = true } tracing = { workspace = true } +unicode-normalization = { workspace = true, optional = true } url = { workspace = true } uuid = { workspace = true, features = ["v7"] } wildmatch = { workspace = true } @@ -113,6 +116,7 @@ xai-test-utils = { workspace = true } [build-dependencies] reqwest = { workspace = true, features = ["blocking", "rustls-tls"] } flate2 = { workspace = true } +sha2 = { workspace = true } tar = { workspace = true } [lints] diff --git a/crates/codegen/xai-grok-tools/build.rs b/crates/codegen/xai-grok-tools/build.rs index e021c0f..8a93533 100644 --- a/crates/codegen/xai-grok-tools/build.rs +++ b/crates/codegen/xai-grok-tools/build.rs @@ -10,15 +10,185 @@ use std::path::PathBuf; const RG_VER: &str = "15.0.0"; const BFS_VER: &str = "4.1"; const UGREP_VER: &str = "7.7.0"; +const FD_VER: &str = "10.4.2"; +// fd stopped publishing x86_64-apple-darwin assets after 10.3.0. +const FD_VER_MACOS_X64: &str = "10.3.0"; + +/// Pinned SHA-256 of each `(version, triple)` fd release tarball we embed. +const FD_TARBALL_SHA256: &[(&str, &str, &str)] = &[ + ( + "10.4.2", + "x86_64-unknown-linux-musl", + "e3257d48e29a6be965187dbd24ce9af564e0fe67b3e73c9bdcd180f4ec11bdde", + ), + ( + "10.4.2", + "aarch64-unknown-linux-musl", + "f32d3657473fba74e2600babc8db0b93420d51169223b7e8143b2ed55d8fd9e8", + ), + ( + "10.4.2", + "aarch64-apple-darwin", + "623dc0afc81b92e4d4606b380d7bc91916ba7b97814263e554d50923a39e480a", + ), + ( + "10.3.0", + "x86_64-apple-darwin", + "50d30f13fe3d5914b14c4fff5abcbd4d0cdab4b855970a6956f4f006c17117a3", + ), +]; fn main() -> Result<(), Box> { bundle_rg()?; + // fd is an optional vendored file-search binary backing a feature-gated + // toolset; skip the download/embed entirely when that feature is off + // (shipped TUI binaries). + if env::var_os("CARGO_FEATURE_PI").is_some() { + bundle_fd()?; + } // bfs/ugrep back the bash-harness find/grep shadows (embedded_search_tools). bundle_search_tool("bfs", "BFS", BFS_VER)?; bundle_search_tool("ugrep", "UGREP", UGREP_VER)?; Ok(()) } +/// Download + embed fd as an optional vendored file-search binary, mirroring +/// the ripgrep bundling +/// (release-only or `GROK_TOOLS_BUNDLE_FD_PATH` override), plus pinned +/// per-asset SHA-256 verification of the downloaded tarball. +fn bundle_fd() -> Result<(), Box> { + println!("cargo:rerun-if-env-changed=GROK_TOOLS_BUNDLE_FD_PATH"); + println!("cargo:rustc-check-cfg=cfg(bundle_fd)"); + + let gen_dir = PathBuf::from(env::var("OUT_DIR")?).join("bundle-fd"); + fs::create_dir_all(&gen_dir)?; + + // The consuming vendor extraction is unix-only — never bundle on + // Windows targets, mirroring the bfs/ugrep skip. + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os == "windows" { + return Ok(()); + } + + let path_override = env::var("GROK_TOOLS_BUNDLE_FD_PATH").ok(); + let is_release = env::var("PROFILE").as_deref() == Ok("release"); + if path_override.is_none() && !is_release { + return Ok(()); + } + + // Per-target version: macOS x86_64 pins the last release with that asset. + let target_arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap_or_default(); + let (ver, asset_triple) = match (target_os.as_str(), target_arch.as_str()) { + ("macos", "aarch64") => (FD_VER, "aarch64-apple-darwin"), + ("macos", "x86_64") => (FD_VER_MACOS_X64, "x86_64-apple-darwin"), + ("linux", "x86_64") => (FD_VER, "x86_64-unknown-linux-musl"), + ("linux", "aarch64") => (FD_VER, "aarch64-unknown-linux-musl"), + _ => { + if path_override.is_none() { + return Err(format!( + "Unsupported target for fd bundling: {target_os}-{target_arch}. Set GROK_TOOLS_BUNDLE_FD_PATH to a local fd binary for offline or unsupported builds.", + ) + .into()); + } + (FD_VER, "override") + } + }; + + println!("cargo:rustc-cfg=bundle_fd"); + println!("cargo:rustc-env=GROK_TOOLS_FD_VER={ver}"); + + if let Some(path) = path_override { + let dest = gen_dir.join(format!("fd-{ver}-override.bin")); + println!("cargo:rustc-env=GROK_TOOLS_FD_TARGET=override"); + let _ = fs::remove_file(&dest); + fs::copy(PathBuf::from(path.clone()), &dest).map_err(|e| { + format!( + "Failed copying GROK_TOOLS_BUNDLE_FD_PATH: {e} from path {path} to dest {}", + dest.display() + ) + })?; + return Ok(()); + } + + println!("cargo:rustc-env=GROK_TOOLS_FD_TARGET={asset_triple}"); + let dest = gen_dir.join(format!("fd-{ver}-{asset_triple}.bin")); + let _ = fs::remove_file(&dest); + + let url = format!( + "https://github.com/sharkdp/fd/releases/download/v{ver}/fd-v{ver}-{asset_triple}.tar.gz" + ); + + let bytes: Vec = { + let resp = reqwest::blocking::get(&url).map_err(|e| { + format!( + "Failed to download fd: {e}\nSet GROK_TOOLS_BUNDLE_FD_PATH to a local fd for offline builds." + ) + })?; + if !resp.status().is_success() { + return Err(format!( + "HTTP {} downloading fd. Set GROK_TOOLS_BUNDLE_FD_PATH for offline builds.", + resp.status() + ) + .into()); + } + resp.bytes()?.to_vec() + }; + + // Verify the tarball against the pinned per-asset hash before unpacking. + let expected_sha = FD_TARBALL_SHA256 + .iter() + .find(|(v, t, _)| *v == ver && *t == asset_triple) + .map(|(_, _, sha)| *sha) + .ok_or_else(|| format!("No pinned SHA-256 for fd {ver} {asset_triple}"))?; + let actual_sha = { + use sha2::Digest as _; + let mut hasher = sha2::Sha256::new(); + hasher.update(&bytes); + hex_encode(&hasher.finalize()) + }; + if actual_sha != expected_sha { + return Err(format!( + "SHA-256 mismatch for {url}:\n expected {expected_sha}\n actual {actual_sha}" + ) + .into()); + } + + let gz = flate2::read::GzDecoder::new(&bytes[..]); + let mut ar = tar::Archive::new(gz); + let mut found = false; + for entry in ar.entries()? { + let mut e = entry?; + let p = e.path()?; + if p.file_name().is_some_and(|n| n == "fd") { + let data: Vec = { + let mut v = Vec::new(); + io::copy(&mut e, &mut v)?; + v + }; + fs::write(&dest, &data)?; + found = true; + break; + } + } + + if !found { + return Err(format!( + "Could not find 'fd' in fd archive {url}. Set GROK_TOOLS_BUNDLE_FD_PATH for offline builds." + ) + .into()); + } + + Ok(()) +} + +fn hex_encode(bytes: &[u8]) -> String { + let mut out = String::with_capacity(bytes.len() * 2); + for byte in bytes { + out.push_str(&format!("{byte:02x}")); + } + out +} + /// Bundle a prebuilt **static** search-tool binary (`bfs`/`ugrep`) when /// `GROK_TOOLS_BUNDLE__PATH` points at one (supplied by the release /// pipeline). Emits diff --git a/crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/mod.rs b/crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/mod.rs index a1a05a0..c467877 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/mod.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/mod.rs @@ -31,29 +31,22 @@ use xai_tool_types::{ /// constant is not applied unless a wait is active. pub(crate) const DEFAULT_WAIT_TIMEOUT: Duration = Duration::from_secs(30); -/// Max time a blocking wait (`get_command_or_subagent_output` with positive -/// `timeout_ms` / `wait_commands_or_subagents`) may hold the turn, regardless of -/// the requested `timeout_ms`. Safe to cap because completed tasks ping the -/// model (`send_task_complete` → auto-wake). 10m matches the external -/// `TaskOutput` cap. Env override: `GROK_MAX_WAIT_BLOCK_MS`. -const MAX_WAIT_BLOCK: Duration = Duration::from_secs(600); - -fn max_wait_block() -> Duration { - std::env::var("GROK_MAX_WAIT_BLOCK_MS") - .ok() - .and_then(|s| s.parse::().ok()) - .map(Duration::from_millis) - .unwrap_or(MAX_WAIT_BLOCK) +/// The blocking-wait ceiling: `GROK_MAX_WAIT_BLOCK_MS`, else 10 min. +/// +/// The same value fills `{max_wait_ms}` in the descriptions, so a wait can +/// never exceed what the model was told it may ask for. +pub(crate) fn max_wait_block() -> Duration { + Duration::from_millis(xai_tool_types::max_wait_block_ms()) } /// Resolve a model-supplied `timeout_ms` into the effective blocking-wait -/// duration: default when omitted, then clamped to [`max_wait_block`] so a -/// single wait call can never wedge the turn for longer than the cap. -pub(crate) fn capped_wait_timeout(timeout_ms: Option) -> Duration { +/// duration: default when omitted, then clamped to `cap` so a single wait call +/// can never wedge the turn for longer than the ceiling. +pub(crate) fn capped_wait_timeout(timeout_ms: Option, cap: Duration) -> Duration { let base = timeout_ms .map(Duration::from_millis) .unwrap_or(DEFAULT_WAIT_TIMEOUT); - base.min(max_wait_block()) + base.min(cap) } /// The caller's requested wait before capping, or the default when omitted. @@ -201,10 +194,11 @@ impl TaskOutputTool { } let waits = xai_tool_types::task_output_waits(timeout_ms); + let wait_cap = max_wait_block(); let wait_hint = if waits { WaitHint::Elapsed { requested: requested_wait_timeout(timeout_ms), - waited: capped_wait_timeout(timeout_ms), + waited: capped_wait_timeout(timeout_ms, wait_cap), } } else { WaitHint::NotRequested @@ -212,7 +206,7 @@ impl TaskOutputTool { let snapshot = if waits { // Cap the blocking wait so a large `timeout_ms` can't wedge the turn; // the model is pinged on completion regardless (see `capped_wait_timeout`). - let timeout = capped_wait_timeout(timeout_ms); + let timeout = capped_wait_timeout(timeout_ms, wait_cap); terminal.wait_for_completion(task_id, Some(timeout)).await } else { terminal.get_task(task_id).await @@ -255,7 +249,7 @@ impl TaskOutputTool { // Same cap as the bash path: a blocking subagent query can't wedge the // turn beyond the wait cap (the parent is pinged when the child finishes). let query_timeout_ms = if waits { - Some(capped_wait_timeout(timeout_ms).as_millis() as u64) + Some(capped_wait_timeout(timeout_ms, wait_cap).as_millis() as u64) } else { timeout_ms }; @@ -298,7 +292,7 @@ impl TaskOutputTool { ) -> Result { let waits = xai_tool_types::task_output_waits(timeout_ms); let requested = requested_wait_timeout(timeout_ms); - let timeout = capped_wait_timeout(timeout_ms); + let timeout = capped_wait_timeout(timeout_ms, max_wait_block()); let (terminal, backend, read_file_name, max_output_bytes) = { let res = resources.lock().await; @@ -1087,13 +1081,26 @@ mod tests { // unbounded blocking wait wedged the turn for hours). #[test] fn capped_wait_timeout_clamps_and_defaults() { - assert_eq!(capped_wait_timeout(None), DEFAULT_WAIT_TIMEOUT); + let cap = Duration::from_millis(xai_tool_types::MAX_WAIT_BLOCK_MS_DEFAULT); + assert_eq!(capped_wait_timeout(None, cap), DEFAULT_WAIT_TIMEOUT); assert_eq!( - capped_wait_timeout(Some(5_000)), + capped_wait_timeout(Some(5_000), cap), Duration::from_millis(5_000) ); - assert_eq!(capped_wait_timeout(Some(36_000_000)), MAX_WAIT_BLOCK); - assert_eq!(capped_wait_timeout(Some(600_000)), MAX_WAIT_BLOCK); + assert_eq!(capped_wait_timeout(Some(36_000_000), cap), cap); + assert_eq!(capped_wait_timeout(Some(600_000), cap), cap); + } + + /// A client that shortens the cap at finalize must also shorten the wait — + /// otherwise the server outlasts the deadline the client will honor. + #[test] + fn capped_wait_timeout_honors_a_shortened_cap() { + let cap = Duration::from_millis(300_000); + assert_eq!(capped_wait_timeout(Some(600_000), cap), cap); + assert_eq!( + capped_wait_timeout(Some(120_000), cap), + Duration::from_millis(120_000) + ); } #[test] diff --git a/crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/terminal_command.rs b/crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/terminal_command.rs index 73deb00..19717c9 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/terminal_command.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/terminal_command.rs @@ -22,11 +22,14 @@ impl crate::types::tool_metadata::ToolMetadata for GetTerminalCommandOutputTool } fn description_template(&self) -> &str { + // `{max_wait_ms}` is resolved per session by the finalize loop's + // `TruncationConfig::interpolate_description`, like `{max_lines_read}`: + // the cap is client-configurable, so it cannot be baked in here. r#"Get output and status from a background terminal command${%- if tools.by_kind.monitor %} or monitor${%- endif %}. Usage notes: - Pass ${{ params.background_task_action.task_ids }} with one or more ids from ${%- if params is defined and params.execute is defined and params.execute.is_background %} ${{ params.execute.is_background }}=true commands${%- else %} background commands${%- endif %}${%- if tools.by_kind.monitor %} (a monitor's ${{ params.kill_task_action.task_id }} is returned by ${{ tools.by_kind.monitor }})${%- endif %}; for a single task use a one-element array. Multiple ids with a positive ${{ params.background_task_action.timeout_ms }} wait until all complete -- Omit ${{ params.background_task_action.timeout_ms }} or pass 0 for a non-blocking status snapshot; set a positive ${{ params.background_task_action.timeout_ms }} to wait up to that many milliseconds, capped at ~10 min +- Omit ${{ params.background_task_action.timeout_ms }} or pass 0 for a non-blocking status snapshot; set a positive ${{ params.background_task_action.timeout_ms }} to wait up to that many milliseconds, capped at {max_wait_ms} - Returns current output, status, and exit code if completed${%- if tools.by_kind.read %} - If output is large, use ${{ tools.by_kind.read }} on the output_file path${%- endif %}"# } diff --git a/crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/wait_tasks.rs b/crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/wait_tasks.rs index 71747b4..be26a98 100644 --- a/crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/wait_tasks.rs +++ b/crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/wait_tasks.rs @@ -173,8 +173,10 @@ impl xai_tool_runtime::Tool for WaitTasksTool { .timeout_ms .map(std::time::Duration::from_millis) .unwrap_or(super::DEFAULT_WAIT_TIMEOUT); - let timeout = - crate::implementations::grok_build::task_output::capped_wait_timeout(input.timeout_ms); + let timeout = crate::implementations::grok_build::task_output::capped_wait_timeout( + input.timeout_ms, + crate::implementations::grok_build::task_output::max_wait_block(), + ); let (terminal, backend, read_file_name, max_output_bytes) = { let res = resources.lock().await; diff --git a/crates/codegen/xai-grok-tools/src/registry/types.rs b/crates/codegen/xai-grok-tools/src/registry/types.rs index b7a96f9..ca843f0 100644 --- a/crates/codegen/xai-grok-tools/src/registry/types.rs +++ b/crates/codegen/xai-grok-tools/src/registry/types.rs @@ -1162,9 +1162,16 @@ impl ToolRegistryBuilder { desc, &client_name, crate::DEFAULT_TOOL_OUTPUT_BYTES, + xai_tool_types::max_wait_block_ms(), )); } renderer.render_schema_descriptions(&mut definition.function.parameters); + truncation_config.apply_to_schema( + &mut definition.function.parameters, + &client_name, + crate::DEFAULT_TOOL_OUTPUT_BYTES, + xai_tool_types::max_wait_block_ms(), + ); (entry.apply_params)(&effective_params, &mut resources); tools.push(FinalizedTool { namespace: entry.namespace, @@ -1323,6 +1330,13 @@ impl FinalizedToolset { pub fn local_registry(&self) -> &xai_computer_hub_sdk::LocalRegistry { &self.local_registry } + /// Whether the server must await this tool's in-process cancellation cleanup. + pub fn cooperative_cancellation(&self, tool_name: &str) -> bool { + { + let _ = tool_name; + false + } + } /// Get all tool definitions to send to the client. pub fn tool_definitions(&self) -> Vec { self.tools @@ -2336,6 +2350,7 @@ mod tests { .into_iter() .map(|id| ToolConfig::from_id(format!("GrokBuild:{id}"))) .chain(std::iter::empty::()) + .chain(std::iter::empty::()) .collect(), behavior_preset: None, }; @@ -2389,7 +2404,8 @@ mod tests { _ => {} } } - for def in toolset.tool_definitions() { + let definitions = toolset.tool_definitions(); + for def in definitions { let name = &def.function.name; let desc = def.function.description.as_deref().unwrap_or_default(); assert!( @@ -2422,6 +2438,10 @@ mod tests { collect_descriptions(&def.function.parameters, &mut field_descs); for field_desc in &field_descs { assert_no_render_whitespace_artifacts(name, field_desc); + assert!( + !field_desc.contains("{max_"), + "{name}: unresolved {{max_*}} placeholder in a field description" + ); } } } @@ -3428,6 +3448,119 @@ mod tests { ); } #[tokio::test] + async fn non_pi_finalized_contract_snapshot_is_unchanged() { + let tmp = TempDir::new().unwrap(); + let toolset = ToolRegistryBuilder::new() + .finalize( + ToolServerConfig { + tools: vec![ + ToolConfig::for_tool::(), + ToolConfig::for_tool::(), + ], + behavior_preset: None, + }, + test_session_context(&tmp), + ) + .unwrap(); + let mut contracts: Vec = toolset + .tool_definitions() + .into_iter() + .map(|definition| { + serde_json::json!({ + "name": definition.function.name, + "description": definition.function.description, + "parameters": definition.function.parameters, + }) + }) + .collect(); + contracts.sort_by(|a, b| a["name"].as_str().cmp(&b["name"].as_str())); + let expected: serde_json::Value = serde_json::from_str( + r##" + [ + { + "name": "todo_write", + "description": "Create and manage a structured task list. The user sees this list live — it is your primary way to show progress.\n\nUse for any task with 3+ steps. Skip for trivial single-step work.", + "parameters": { + "$schema": "http://json-schema.org/draft-07/schema#", + "required": [ + "todos" + ], + "type": "object", + "properties": { + "merge": { + "description": "Optional. When true (default), merges the provided todos into the existing list by id — send only the items you are changing, and to flip status without changing content send just id + status. When false, the provided todos replace the existing list.", + "type": "boolean", + "default": true + }, + "todos": { + "description": "Array of todo items to write to the workspace", + "type": "array", + "items": { + "type": "object", + "properties": { + "id": { + "description": "Unique identifier for the todo item", + "type": "string" + }, + "content": { + "description": "The description/content of the todo item", + "type": [ + "string", + "null" + ] + }, + "status": { + "description": "The status of the todo item: pending, in_progress, completed, or cancelled", + "type": [ + "string", + "null" + ], + "enum": [ + "pending", + "in_progress", + "completed", + "cancelled", + null + ] + } + }, + "required": [ + "id" + ] + } + } + } + } + }, + { + "name": "write", + "description": "Create or overwrite a file.\n\n- Writing to an existing path replaces the file.\n- Parent directories are created for you.", + "parameters": { + "$schema": "http://json-schema.org/draft-07/schema#", + "required": [ + "file_path", + "content" + ], + "properties": { + "file_path": { + "description": "The absolute path to the file to write.", + "type": "string" + }, + "content": { + "description": "The full file content to write.", + "type": "string" + } + }, + "type": "object" + } + } + ] +"##, + ) + .expect("checked-in snapshot parses"); + assert_eq!(expected, serde_json::Value::Array(contracts)); + } + #[tokio::test] async fn tool_definitions_builtins_only_hides_mcp_tools() { let tmp = TempDir::new().unwrap(); let builder = ToolRegistryBuilder::new(); @@ -3805,6 +3938,22 @@ mod tests { desc.contains("run_in_background"), "`{name}` description must resolve params.task.run_in_background" ); + let timeout = defs + .iter() + .find(|d| d.function.name == name) + .map(|d| &d.function.parameters["properties"]["timeout_ms"]) + .unwrap_or_else(|| panic!("`{name}` should expose timeout_ms")); + assert!( + timeout.get("maximum").is_some(), + "`{name}`.timeout_ms must carry the resolved wait ceiling: {timeout}" + ); + assert!( + !timeout["description"] + .as_str() + .unwrap_or("") + .contains("{max_"), + "`{name}`.timeout_ms has an unresolved placeholder: {timeout}" + ); } } #[test] diff --git a/crates/codegen/xai-grok-tools/src/reminders/task_completion.rs b/crates/codegen/xai-grok-tools/src/reminders/task_completion.rs index 4c76077..aad8f1b 100644 --- a/crates/codegen/xai-grok-tools/src/reminders/task_completion.rs +++ b/crates/codegen/xai-grok-tools/src/reminders/task_completion.rs @@ -402,6 +402,10 @@ pub async fn resolve_read_tool_name(bridge: &ToolBridge) -> Option { /// in the current agent's toolset. When `None`, the subagent's full /// `output` is inlined verbatim -- this notification is the only place /// the model will see it (no disk-backed output file exists for subagents). +/// +/// KEEP IN SYNC: the exact wording of this message is a compatibility +/// surface — downstream mirrors reproduce it verbatim (grep for +/// `format_subagent_completion_reminder`). Update them when changing it. pub fn format_subagent_completion( c: &SubagentCompletionSummary, task_output_name: Option<&str>, diff --git a/crates/codegen/xai-grok-tools/src/types/context.rs b/crates/codegen/xai-grok-tools/src/types/context.rs index ebdde58..c9ee2fb 100644 --- a/crates/codegen/xai-grok-tools/src/types/context.rs +++ b/crates/codegen/xai-grok-tools/src/types/context.rs @@ -66,6 +66,7 @@ impl TruncationConfig { /// /// Recognized placeholders: /// - `{max_lines_read}` — from `max_lines_read` (default 1000) + /// - `{max_wait_ms}` — the blocking-wait ceiling, as `600000 (~10 min)` /// - `{max_output_bytes}` — resolved via `max_output_bytes_for(tool_name, builtin_default)` /// - `{max_chars_per_line}` — fixed display value for opencode-compat /// descriptions only; the opencode `read` tool clips at its own @@ -78,9 +79,14 @@ impl TruncationConfig { description: &str, tool_name: &str, builtin_output_default: usize, + max_wait_ms: u64, ) -> String { description .replace("{max_lines_read}", &self.max_lines_read().to_string()) + .replace( + "{max_wait_ms}", + &xai_tool_types::format_wait_cap_ms(max_wait_ms), + ) .replace("{max_chars_per_line}", "2000") .replace( "{max_output_bytes}", @@ -89,6 +95,50 @@ impl TruncationConfig { .to_string(), ) } + + /// Resolve placeholders in each schema property description, and pin the + /// blocking-wait ceiling as a `maximum` on whichever property documents it. + /// + /// The tool description alone cannot carry the cap: `description_override` + /// replaces that string outright under toolchain randomization, so on most + /// draws the interpolated copy never reaches the model. Properties are only + /// ever renamed, so a bound placed here survives every draw — and a + /// `maximum` reaches a model that skips the prose. + /// + /// `{max_wait_ms}` in a property description is the marker for which + /// property is the wait, so no tool or parameter name is hardcoded and a + /// renamed parameter is handled for free (keys are remapped by the time + /// this runs). + pub fn apply_to_schema( + &self, + schema: &mut serde_json::Value, + tool_name: &str, + builtin_output_default: usize, + max_wait_ms: u64, + ) { + let Some(properties) = schema.get_mut("properties").and_then(|p| p.as_object_mut()) else { + return; + }; + for property in properties.values_mut() { + let Some(object) = property.as_object_mut() else { + continue; + }; + let Some(description) = object.get("description").and_then(|d| d.as_str()) else { + continue; + }; + let documents_wait = description.contains("{max_wait_ms}"); + let resolved = self.interpolate_description( + description, + tool_name, + builtin_output_default, + max_wait_ms, + ); + object.insert("description".to_string(), serde_json::json!(resolved)); + if documents_wait { + object.insert("maximum".to_string(), serde_json::json!(max_wait_ms)); + } + } + } } #[cfg(test)] @@ -108,6 +158,114 @@ mod tests { assert_eq!(cfg.max_lines_read(), 50); } + #[test] + fn interpolate_description_resolves_max_wait_ms() { + let cfg = TruncationConfig::default(); + let cap = 300_000; + assert_eq!( + cfg.interpolate_description("capped at {max_wait_ms}", "get_task_output", 40_000, cap), + "capped at 300000 (~5 min)" + ); + let default_cap = + xai_tool_types::format_wait_cap_ms(xai_tool_types::MAX_WAIT_BLOCK_MS_DEFAULT); + assert_eq!( + TruncationConfig::default().interpolate_description( + "capped at {max_wait_ms}", + "get_task_output", + 40_000, + xai_tool_types::MAX_WAIT_BLOCK_MS_DEFAULT, + ), + format!("capped at {default_cap}") + ); + } + + #[test] + fn apply_to_schema_resolves_and_pins_the_wait_property() { + let cfg = TruncationConfig::default(); + let cap = 300_000; + let mut schema = serde_json::json!({ + "properties": { + "timeout_ms": {"type": "integer", "description": "Wait up to {max_wait_ms}."}, + "task_ids": {"type": "array", "description": "Task IDs."}, + } + }); + cfg.apply_to_schema(&mut schema, "get_task_output", 40_000, cap); + + let timeout = &schema["properties"]["timeout_ms"]; + assert_eq!(timeout["description"], "Wait up to 300000 (~5 min)."); + assert_eq!(timeout["maximum"], serde_json::json!(300_000u64)); + // Only the property documenting the wait gets a ceiling. + assert_eq!(schema["properties"]["task_ids"]["description"], "Task IDs."); + assert!(schema["properties"]["task_ids"].get("maximum").is_none()); + } + + #[test] + fn apply_to_schema_tracks_a_raised_ceiling_and_a_renamed_property() { + // A 900s actor must not be handed the 300s default, and the marker — + // not the property name — is what identifies the wait. + let cfg = TruncationConfig::default(); + let cap = 900_000; + let mut schema = serde_json::json!({ + "properties": { + "max_wait": {"type": "integer", "description": "Up to {max_wait_ms}."}, + } + }); + cfg.apply_to_schema(&mut schema, "get_task_output", 40_000, cap); + + assert_eq!( + schema["properties"]["max_wait"]["description"], + "Up to 900000 (~15 min)." + ); + assert_eq!( + schema["properties"]["max_wait"]["maximum"], + serde_json::json!(900_000u64) + ); + } + + /// The bound has to land somewhere that actually constrains the value. + /// `Option` could plausibly be emitted as `anyOf: [integer, null]`, in + /// which case a root `maximum` would be inert — so assert against the real + /// generated schema rather than a hand-written one, and pin the shape it + /// relies on. schemars puts `minimum` at the root for the same field, which + /// is the precedent this follows. + #[test] + fn apply_to_schema_bounds_the_real_optional_u64_property() { + let generated = + serde_json::to_value(schemars::schema_for!(xai_tool_types::TaskOutputToolInput)) + .unwrap(); + let timeout = &generated["properties"]["timeout_ms"]; + assert!( + timeout.get("anyOf").is_none(), + "shape changed to anyOf — a root `maximum` no longer constrains the \ + integer arm, so apply_to_schema must walk the branches: {timeout}" + ); + assert_eq!(timeout["type"], serde_json::json!(["integer", "null"])); + + let cfg = TruncationConfig::default(); + let cap = 300_000; + let mut schema = generated.clone(); + cfg.apply_to_schema(&mut schema, "get_task_output", 40_000, cap); + + let bounded = &schema["properties"]["timeout_ms"]; + assert_eq!(bounded["maximum"], serde_json::json!(300_000u64)); + assert!( + !bounded["description"].as_str().unwrap().contains("{max_"), + "placeholder survived: {bounded}" + ); + } + + #[test] + fn apply_to_schema_tolerates_schemas_without_properties() { + let mut schema = serde_json::json!({"type": "object"}); + TruncationConfig::default().apply_to_schema( + &mut schema, + "get_task_output", + 40_000, + xai_tool_types::MAX_WAIT_BLOCK_MS_DEFAULT, + ); + assert_eq!(schema, serde_json::json!({"type": "object"})); + } + #[test] fn mcp_max_output_bytes_for_lookup_order() { // per-tool > mcp-specific > default > builtin diff --git a/crates/codegen/xai-grok-version/Cargo.toml b/crates/codegen/xai-grok-version/Cargo.toml index 2125e82..c7c3e43 100644 --- a/crates/codegen/xai-grok-version/Cargo.toml +++ b/crates/codegen/xai-grok-version/Cargo.toml @@ -1,7 +1,7 @@ [package] license = "Apache-2.0" name = "xai-grok-version" -version = "0.2.116" +version = "0.2.117" edition.workspace = true description = "Lockstepped grok CLI version." diff --git a/crates/codegen/xai-grok-workspace-types/src/rpc/export.rs b/crates/codegen/xai-grok-workspace-types/src/rpc/export.rs new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/crates/codegen/xai-grok-workspace-types/src/rpc/export.rs @@ -0,0 +1 @@ + diff --git a/crates/codegen/xai-grok-workspace-types/src/rpc/mod.rs b/crates/codegen/xai-grok-workspace-types/src/rpc/mod.rs index 378beb0..3fb0886 100644 --- a/crates/codegen/xai-grok-workspace-types/src/rpc/mod.rs +++ b/crates/codegen/xai-grok-workspace-types/src/rpc/mod.rs @@ -13,6 +13,7 @@ pub mod agents_md; pub mod code_nav; pub mod deploy; pub mod envelope; +pub mod export; pub mod export_github; pub mod fs; pub mod git; diff --git a/crates/codegen/xai-grok-workspace/Cargo.toml b/crates/codegen/xai-grok-workspace/Cargo.toml index e48a677..b3e202d 100644 --- a/crates/codegen/xai-grok-workspace/Cargo.toml +++ b/crates/codegen/xai-grok-workspace/Cargo.toml @@ -92,6 +92,8 @@ urlencoding = "2" xai-fast-worktree = { path = "../xai-fast-worktree", features = ["metadata"] } # tonic: Status/Code mapping for deploy errors in workspace_ops. tonic = { workspace = true } +sha1 = { workspace = true, optional = true } +zip = { workspace = true, optional = true } xai-fsnotify = { path = "../xai-fsnotify" } clap = { workspace = true } diff --git a/crates/codegen/xai-grok-workspace/src/activity.rs b/crates/codegen/xai-grok-workspace/src/activity.rs index f1bfcbf..f0b5c75 100644 --- a/crates/codegen/xai-grok-workspace/src/activity.rs +++ b/crates/codegen/xai-grok-workspace/src/activity.rs @@ -8,7 +8,7 @@ use std::time::Instant; use dashmap::DashMap; use xai_file_utils::events::{Event, EventWriter, ToolCompletedSource, ToolOutcome}; use xai_file_utils::queue::UploadQueueStats; -use xai_tool_protocol::{ToolServerLifecycleStatus, ToolServerStatusPayload}; +use xai_tool_protocol::{IdleWithholdReason, ToolServerLifecycleStatus, ToolServerStatusPayload}; const LIFECYCLE_NONE: u8 = 0; const LIFECYCLE_DRAINING: u8 = 1; @@ -88,10 +88,25 @@ pub struct ActivityTracker { /// Window (ms) recent preview-proxy traffic withholds idle for; defaults to /// [`PREVIEW_ACTIVITY_WINDOW_MS`], overridable via the builder. preview_activity_window_ms: u64, - /// Epoch ms of the last scraped preview-proxy activity (`0` = none). Fed by - /// the preview-activity scraper (`preview_supervisor`); withholds idle within + /// Epoch ms the pane's own status poll was last observed (`0` = none). Fed + /// by the preview-activity scraper; withholds idle within /// [`preview_activity_window_ms`](Self::preview_activity_window_ms). - last_preview_activity_ms: AtomicU64, + /// + /// The *observation* time, not the proxy's stamp: the two processes are not + /// clock-coupled, and only the local clock is comparable with the rest. + last_preview_status_ms: AtomicU64, + /// Epoch ms real app traffic was last observed (`0` = none). + last_preview_routed_ms: AtomicU64, + /// Open preview WebSocket (HMR) tunnels as of the last scrape. Nonzero ⇒ a + /// client is attached, which no activity stamp would reveal. + preview_ws_tunnels_open: AtomicU64, + /// In-flight `Routed` preview requests as of the last scrape. + preview_routed_in_flight: AtomicU64, + /// Epoch ms this process started — the floor of the withhold anchor, so a + /// young or freshly-restored workspace is never treated as long-idle. + /// Distinct from [`Self::started_at`], a monotonic `Instant` that cannot be + /// compared against the epoch stamps around it. + started_at_ms: u64, sessions: DashMap, /// call_id → session_id so `tool_call_completed` can decrement @@ -168,7 +183,11 @@ impl ActivityTracker { durability_idle_hold_max_ms, idle_ignores_background: false, preview_activity_window_ms: PREVIEW_ACTIVITY_WINDOW_MS, - last_preview_activity_ms: AtomicU64::new(0), + last_preview_status_ms: AtomicU64::new(0), + last_preview_routed_ms: AtomicU64::new(0), + preview_ws_tunnels_open: AtomicU64::new(0), + preview_routed_in_flight: AtomicU64::new(0), + started_at_ms: now_ms(), sessions: DashMap::new(), call_to_session: DashMap::new(), prune_window_ms: prune_window.as_millis() as u64, @@ -217,15 +236,54 @@ impl ActivityTracker { self.notify.clone() } - /// Record fresh preview-proxy traffic: withholds `idle_since_ms` for + /// Record fresh `Routed` preview traffic. Withholds `idle_since_ms` for /// [`preview_activity_window_ms`](Self::preview_activity_window_ms) and wakes /// the status publisher so the renewed "active" status reaches the server promptly. - pub fn note_preview_activity(&self) { - self.last_preview_activity_ms + pub fn note_preview_routed_activity(&self) { + self.last_preview_routed_ms .store(now_ms(), Ordering::Relaxed); self.notify.notify_waiters(); } + /// Record a fresh preview status poll. Withholds idle exactly as routed + /// traffic does today, but is tracked separately: the poll continues at the + /// same cadence whether or not anyone is watching. + pub fn note_preview_status_activity(&self) { + self.last_preview_status_ms + .store(now_ms(), Ordering::Relaxed); + self.notify.notify_waiters(); + } + + /// Mirror the proxy's attached-client counters. Absolute values, not edges, + /// so a missed scrape self-corrects on the next one. + pub fn set_preview_attached(&self, ws_tunnels_open: u64, routed_in_flight: u64) { + let was_attached = self.has_preview_client_attached(); + self.preview_ws_tunnels_open + .store(ws_tunnels_open, Ordering::Relaxed); + self.preview_routed_in_flight + .store(routed_in_flight, Ordering::Relaxed); + // The scraper calls this every tick; the common case is 0 → 0. + if was_attached != self.has_preview_client_attached() { + self.notify.notify_waiters(); + } + } + + /// An open WebSocket tunnel or a routed request in flight. Never a poll. + fn has_preview_client_attached(&self) -> bool { + self.preview_ws_tunnels_open.load(Ordering::Relaxed) > 0 + || self.preview_routed_in_flight.load(Ordering::Relaxed) > 0 + } + + pub fn preview_ws_tunnels_open(&self) -> u64 { + self.preview_ws_tunnels_open.load(Ordering::Relaxed) + } + + /// Window recent preview activity withholds idle for — the horizon the + /// stamps decay over, and the one the scraper ages an absent proxy against. + pub fn preview_activity_window_ms(&self) -> u64 { + self.preview_activity_window_ms + } + /// Pending upload-queue items (0 when no queue is coupled). fn upload_queue_pending(&self) -> u64 { self.upload_queue_stats @@ -265,15 +323,33 @@ impl ActivityTracker { let (queue_pending, queue_pending_bytes, queue_inflight, breaker, drain_started) = self.drain_status_fields(); let (producers, durability_withhold) = self.durability_gate(queue_pending, breaker); - // Withhold idle on durability work OR recent preview traffic, decided here + let now = now_ms(); + let (preview_withhold, preview_reason, preview_anchor) = self.preview_withholds_idle(now); + // Withhold idle on durability work OR preview activity, decided here // once so both snapshot paths agree (preview has no hold cap; 12h VM TTL backstops). - let withhold_idle = durability_withhold || self.preview_withholds_idle(now_ms()); + let withhold_idle = durability_withhold || preview_withhold; + // Invariants: no reason while genuinely busy (`idle_since == 0` — the + // work is the cause, not a concurrent poll); durability outranks + // preview; every reason carries a stamp. + let (withhold_reason, withhold_since_ms) = if idle_since == 0 { + (None, None) + } else if durability_withhold { + ( + Some(IdleWithholdReason::Durability), + Some(self.durability_busy_since_ms.load(Ordering::Relaxed)), + ) + } else { + (preview_reason, preview_reason.map(|_| preview_anchor)) + }; DurabilityPayloadFields { idle_since_ms: if idle_since == 0 || withhold_idle { None } else { Some(idle_since) }, + withhold_reason, + withhold_since_ms, + preview_ws_tunnels_open: self.preview_ws_tunnels_open().min(u32::MAX as u64) as u32, upload_queue_pending: queue_pending, upload_queue_pending_bytes: queue_pending_bytes, upload_queue_inflight: queue_inflight, @@ -316,13 +392,41 @@ impl ActivityTracker { (producers, !hold_expired) } - /// Whether recent preview-proxy traffic should currently withhold idle. - fn preview_withholds_idle(&self, now: u64) -> bool { - preview_activity_withholds_idle( + /// Whether preview activity should withhold idle, and on what grounds. + /// + /// Returns `(withhold, reason, anchor)`, where the anchor is the epoch-ms + /// the current hold is measured from. Tiers are checked strongest first; + /// all three withhold identically today, only the accounting differs. + fn preview_withholds_idle(&self, now: u64) -> (bool, Option, u64) { + // Including process start means a young or freshly-restored workspace + // can never look long-idle — the process restarts on restore, so this + // covers revived sessions without a separate minimum-age rule. + let anchor = self + .last_preview_routed_ms + .load(Ordering::Relaxed) + .max(self.last_call_completed_ms.load(Ordering::Relaxed)) + .max(self.started_at_ms); + + if self.has_preview_client_attached() { + return (true, Some(IdleWithholdReason::PreviewAttached), anchor); + } + if preview_activity_withholds_idle( now, - self.last_preview_activity_ms.load(Ordering::Relaxed), + self.last_preview_routed_ms.load(Ordering::Relaxed), self.preview_activity_window_ms, - ) + ) { + return (true, Some(IdleWithholdReason::PreviewRouted), anchor); + } + if preview_activity_withholds_idle( + now, + self.last_preview_status_ms.load(Ordering::Relaxed), + self.preview_activity_window_ms, + ) { + // Holds, but never advances the anchor: a poll must not reset a + // clock meant to measure real use. + return (true, Some(IdleWithholdReason::PreviewStatusOnly), anchor); + } + (false, None, anchor) } /// Whether any tracked session currently has an active turn (the aggregate @@ -522,12 +626,16 @@ impl ActivityTracker { count } - /// Mark a session as ended: clear turn-active flag and notify waiters. - /// - /// Called by [`crate::handle::WorkspaceHandle::on_session_ended()`] when - /// a `HookEvent::SessionEnded` arrives from the server. + /// Releases the session's entry. One with calls in flight keeps it: the + /// swap policy reads that count, and the idle prune collects it later. pub fn session_ended(&self, session_id: &str) { - if let Some(session) = self.sessions.get(session_id) { + let released = self + .sessions + .remove_if(session_id, |_, s| { + s.active_tool_calls.load(Ordering::Acquire) == 0 + }) + .is_some(); + if !released && let Some(session) = self.sessions.get(session_id) { session.turn_active.store(false, Ordering::Release); } self.notify.notify_waiters(); @@ -640,6 +748,11 @@ impl ActivityTracker { } } + /// Resident session records. Does not prune, unlike [`Self::known_sessions`]. + pub fn session_count(&self) -> usize { + self.sessions.len() + } + /// Returns live session IDs. As a side-effect, prunes sessions /// that have been idle longer than the configured prune window. pub fn known_sessions(&self) -> Vec { @@ -731,6 +844,11 @@ impl ActivityTracker { drain_started_ms: d.drain_started_ms, turn_active, idle_ignores_background: self.idle_ignores_background, + withhold_reason: d.withhold_reason, + withhold_since_ms: d.withhold_since_ms, + // No ceilings configured yet, so a hold can never be capped. + withhold_capped: false, + preview_ws_tunnels_open: d.preview_ws_tunnels_open, } } @@ -783,6 +901,11 @@ impl ActivityTracker { drain_started_ms: d.drain_started_ms, turn_active: self.any_turn_active(), idle_ignores_background: self.idle_ignores_background, + withhold_reason: d.withhold_reason, + withhold_since_ms: d.withhold_since_ms, + // No ceilings configured yet, so a hold can never be capped. + withhold_capped: false, + preview_ws_tunnels_open: d.preview_ws_tunnels_open, } } } @@ -791,6 +914,9 @@ impl ActivityTracker { /// [`ActivityTracker::durability_payload_fields`]. struct DurabilityPayloadFields { idle_since_ms: Option, + withhold_reason: Option, + withhold_since_ms: Option, + preview_ws_tunnels_open: u32, upload_queue_pending: u32, upload_queue_pending_bytes: u64, upload_queue_inflight: u32, @@ -1287,29 +1413,195 @@ mod tests { #[test] fn note_preview_activity_withholds_then_resumes_idle() { + for (label, note) in [ + ( + "routed", + &ActivityTracker::note_preview_routed_activity as &dyn Fn(&ActivityTracker), + ), + ("status", &ActivityTracker::note_preview_status_activity), + ] { + let t = ActivityTracker::new(); + assert!( + t.snapshot().idle_since_ms.is_some(), + "{label}: an idle tracker reports idle before any preview activity" + ); + + note(&t); + assert!( + t.snapshot().idle_since_ms.is_none(), + "{label}: recent preview activity must withhold idle" + ); + assert!( + t.snapshot_session("any").idle_since_ms.is_none(), + "{label}: the per-session payload must withhold idle too" + ); + + let stale = now_ms().saturating_sub(PREVIEW_ACTIVITY_WINDOW_MS + 1_000); + t.last_preview_routed_ms.store(stale, Ordering::Relaxed); + t.last_preview_status_ms.store(stale, Ordering::Relaxed); + assert!( + t.snapshot().idle_since_ms.is_some(), + "{label}: idle must resume once the preview window decays" + ); + } + } + + #[test] + fn withhold_reason_reports_the_tier_that_is_holding() { let t = ActivityTracker::new(); - assert!( - t.snapshot().idle_since_ms.is_some(), - "an idle tracker reports idle before any preview activity" + assert_eq!( + t.snapshot().withhold_reason, + None, + "nothing holding ⇒ no reason" ); - t.note_preview_activity(); + t.note_preview_status_activity(); + let s = t.snapshot(); + assert_eq!( + s.withhold_reason, + Some(IdleWithholdReason::PreviewStatusOnly), + "a bare status poll is the weakest tier" + ); + assert!( + s.withhold_since_ms.is_some(), + "every reason carries a since-stamp so a reader can age the hold" + ); + + t.note_preview_routed_activity(); + assert_eq!( + t.snapshot().withhold_reason, + Some(IdleWithholdReason::PreviewRouted), + "real app traffic outranks the status poll" + ); + + t.set_preview_attached(1, 0); + let s = t.snapshot(); + assert_eq!( + s.withhold_reason, + Some(IdleWithholdReason::PreviewAttached), + "an open tunnel outranks everything below it" + ); + assert_eq!(s.preview_ws_tunnels_open, 1); + assert!(!s.withhold_capped, "no ceilings configured yet"); + } + + /// An HMR socket writes no activity stamps, so the counter must hold alone + /// — which is exactly why the scraper must not clear it on one missed poll. + #[test] + fn attached_client_withholds_without_any_activity_stamp() { + let t = ActivityTracker::new(); + assert!(t.snapshot().idle_since_ms.is_some()); + + t.set_preview_attached(1, 0); assert!( t.snapshot().idle_since_ms.is_none(), - "recent preview activity must withhold idle" - ); - assert!( - t.snapshot_session("any").idle_since_ms.is_none(), - "the per-session payload must withhold idle too" + "an open WebSocket tunnel withholds idle by itself" ); - t.last_preview_activity_ms.store( - now_ms().saturating_sub(PREVIEW_ACTIVITY_WINDOW_MS + 1_000), - Ordering::Relaxed, + t.set_preview_attached(0, 1); + assert_eq!( + t.snapshot().withhold_reason, + Some(IdleWithholdReason::PreviewAttached), + "a routed request in flight is equally an attached client" ); + + t.set_preview_attached(0, 0); assert!( t.snapshot().idle_since_ms.is_some(), - "idle must resume once the preview window decays" + "once detached with no stamp in window, idle resumes" + ); + } + + /// A ceiling will be measured from the anchor, so letting the pane's own + /// poll reset it would make that ceiling unreachable. + #[test] + fn status_poll_does_not_advance_the_withhold_anchor() { + let t = ActivityTracker::new(); + t.note_preview_routed_activity(); + let anchor = t.snapshot().withhold_since_ms.expect("routed holds"); + + std::thread::sleep(std::time::Duration::from_millis(5)); + t.note_preview_status_activity(); + assert_eq!( + t.snapshot().withhold_since_ms, + Some(anchor), + "a status poll leaves the anchor where the last real use put it" + ); + } + + /// A session busy with real work must not be attributed to the preview just + /// because the pane happens to be polling it. Both look like a missing + /// `idle_since_ms` on the wire, and conflating them would inflate the + /// preview share of exactly the population this field exists to measure. + #[test] + fn a_genuinely_busy_session_reports_no_withhold_reason() { + let t = ActivityTracker::new(); + t.note_preview_status_activity(); + assert_eq!( + t.snapshot().withhold_reason, + Some(IdleWithholdReason::PreviewStatusOnly), + "idle but polled ⇒ the poll is the reason" + ); + + t.tool_call_started("c1", "read_file", Some("sess-a")); + let s = t.snapshot(); + assert!( + s.idle_since_ms.is_none(), + "a tool call in flight withholds idle on its own" + ); + assert_eq!( + s.withhold_reason, None, + "the tool call is the cause, not the concurrent poll" + ); + assert_eq!(s.withhold_since_ms, None); + + t.tool_call_completed("c1", Some("sess-a"), ToolOutcome::Success); + assert_eq!( + t.snapshot().withhold_reason, + Some(IdleWithholdReason::PreviewStatusOnly), + "once the real work finishes, the poll is the reason again" + ); + } + + /// Durability is already bounded by its own cap, so it is the reason worth + /// reporting when both hold. + #[tokio::test] + async fn durability_outranks_preview_in_the_reported_reason() { + let t = ActivityTracker::new(); + let tasks = tokio_util::task::TaskTracker::new(); + t.set_producer_tasks(tasks.clone()); + + t.note_preview_status_activity(); + assert_eq!( + t.snapshot().withhold_reason, + Some(IdleWithholdReason::PreviewStatusOnly), + "preview alone reports the preview reason" + ); + + // An in-flight producer engages the durability gate. Nothing else does + // — a background task is not durable work. + let gate = Arc::new(tokio::sync::Notify::new()); + let gate2 = gate.clone(); + let join = tasks.spawn(async move { gate2.notified().await }); + + let s = t.snapshot(); + assert_eq!(s.artifact_producers_inflight, 1, "the gate is engaged"); + assert_eq!( + s.withhold_reason, + Some(IdleWithholdReason::Durability), + "durability outranks a concurrent preview hold" + ); + assert!( + s.withhold_since_ms.is_some(), + "a durability hold reports its own busy-since stamp" + ); + + gate.notify_one(); + join.await.expect("producer task must not panic"); + assert_eq!( + t.snapshot().withhold_reason, + Some(IdleWithholdReason::PreviewStatusOnly), + "once durability clears, the preview hold is reported again" ); } @@ -1317,13 +1609,13 @@ mod tests { fn configured_preview_window_overrides_default() { let configured = ActivityTracker::new().with_preview_activity_window_ms(500); configured - .last_preview_activity_ms + .last_preview_status_ms .store(now_ms().saturating_sub(1_000), Ordering::Relaxed); assert!(configured.snapshot().idle_since_ms.is_some()); let default = ActivityTracker::new(); default - .last_preview_activity_ms + .last_preview_status_ms .store(now_ms().saturating_sub(1_000), Ordering::Relaxed); assert!(default.snapshot().idle_since_ms.is_none()); } @@ -1332,7 +1624,7 @@ mod tests { fn preview_activity_does_not_override_active_tool_call() { let t = ActivityTracker::new(); t.tool_call_started("c1", "read_file", Some("sess-a")); - t.note_preview_activity(); + t.note_preview_routed_activity(); let s = t.snapshot(); assert!(s.idle_since_ms.is_none()); assert_eq!( @@ -1704,16 +1996,16 @@ mod tests { } #[test] - fn session_ended_clears_turn_active() { + fn session_ended_releases_an_idle_session() { let t = ActivityTracker::new(); t.turn_started("sess-a", 3); - let session = t.sessions.get("sess-a").expect("session should exist"); - assert!(session.turn_active.load(Ordering::Acquire)); + assert!(t.is_turn_active("sess-a")); t.session_ended("sess-a"); + assert!( - !session.turn_active.load(Ordering::Acquire), - "turn_active should be cleared after session_ended" + !t.sessions.contains_key("sess-a"), + "an ended session must not stay resident" ); } @@ -1732,10 +2024,9 @@ mod tests { t.tool_call_started("c1", "read_file", Some("sess-a")); t.tool_call_completed("c1", None, ToolOutcome::Success); - // session_ended should not panic when turn was never active. t.session_ended("sess-a"); - let session = t.sessions.get("sess-a").expect("session should exist"); - assert!(!session.turn_active.load(Ordering::Acquire)); + + assert!(!t.sessions.contains_key("sess-a")); } #[tokio::test] @@ -1770,8 +2061,11 @@ mod tests { t.session_ended("sess-a"); - // turn_active cleared, but the in-flight tool call remains. assert!(!t.is_turn_active("sess-a")); + assert!( + t.sessions.contains_key("sess-a"), + "a session with a call in flight keeps its counters" + ); assert_eq!( t.snapshot_session("sess-a").active_tool_calls, 1, diff --git a/crates/codegen/xai-grok-workspace/src/error.rs b/crates/codegen/xai-grok-workspace/src/error.rs index 0cc2ad7..4800ad3 100644 --- a/crates/codegen/xai-grok-workspace/src/error.rs +++ b/crates/codegen/xai-grok-workspace/src/error.rs @@ -41,6 +41,10 @@ pub enum WorkspaceError { /// An error from the server connection or tool server. #[error("hub error: {0}")] HubError(String), + #[error("unknown workspace method: {0}")] + UnknownMethod(String), + #[error("workspace archive export failed: {0}")] + ExportArchiveLimitExceeded(String), #[error("github export error: {message}")] ExportGithub { kind: xai_grok_workspace_types::rpc::export_github::ExportGithubError, @@ -78,6 +82,8 @@ impl WorkspaceError { Self::InvalidHunkAction(_) => "invalid_hunk_action", Self::HunkActionFailed(_) => "hunk_action_failed", Self::HubError(_) => "hub_error", + Self::UnknownMethod(_) => "unknown_method", + Self::ExportArchiveLimitExceeded(_) => "export_archive_limit_exceeded", Self::ExportGithub { kind, .. } => kind.wire_code(), Self::ShuttingDown => "shutting_down", Self::ToolsetExternallyOwned(_) => "toolset_externally_owned", diff --git a/crates/codegen/xai-grok-workspace/src/handle.rs b/crates/codegen/xai-grok-workspace/src/handle.rs index 70c891f..816b0c1 100644 --- a/crates/codegen/xai-grok-workspace/src/handle.rs +++ b/crates/codegen/xai-grok-workspace/src/handle.rs @@ -3505,16 +3505,30 @@ impl WorkspaceHandle { let mut any_attempt = false; let mut any_success = false; let session_ids = tracker_for_status.known_sessions(); - for sid in &session_ids { + let mut publish = session_ids.clone(); + for sid in last_sent.keys().filter_map(|k| k.as_ref()) { + if !publish.contains(sid) { + publish.push(sid.clone()); + } + } + let mut closed: Vec = Vec::new(); + for sid in &publish { let payload = tracker_for_status.snapshot_session(sid); let key = Some(sid.clone()); + let ended = !session_ids.iter().any(|s| s == sid); if last_sent.get(&key).map(dedup_key) == Some(dedup_key(&payload)) { + if ended { + closed.push(sid.clone()); + } continue; } if let Some(ok) = send_status(&server_conn, payload.clone()).await { any_attempt = true; if ok { any_success = true; + if ended { + closed.push(sid.clone()); + } last_sent.insert(key, payload); last_successful_send = std::time::Instant::now(); } @@ -3522,7 +3536,10 @@ impl WorkspaceHandle { } last_sent.retain(|k, _| match k { None => true, - Some(sid) => session_ids.iter().any(|s| s == sid), + Some(sid) => { + session_ids.iter().any(|s| s == sid) + || (any_success && !closed.contains(sid)) + } }); let payload = tracker_for_status.snapshot(); let needs_send = last_sent.get(&None).map(dedup_key) != Some(dedup_key(&payload)); diff --git a/crates/codegen/xai-grok-workspace/src/hub_server.rs b/crates/codegen/xai-grok-workspace/src/hub_server.rs index 4023733..9ddf606 100644 --- a/crates/codegen/xai-grok-workspace/src/hub_server.rs +++ b/crates/codegen/xai-grok-workspace/src/hub_server.rs @@ -90,17 +90,13 @@ static WORKSPACE_RPC_DURATION_SECONDS: std::sync::LazyLock = .unwrap() }); const UNKNOWN_METHOD_LABEL: &str = "unknown"; -/// Prefix of the [`WorkspaceError::HubError`] for an unrecognized method. Shared -/// by the dispatch default arm and the metric classifier so the "collapse to -/// `unknown`" decision cannot drift from the error it keys on. -const UNKNOWN_METHOD_ERR_PREFIX: &str = "unknown workspace method:"; /// Zero-init this module's metric families. See [`crate::init_metrics`]. pub(crate) fn init_metrics() { WORKSPACE_RPC_REQUESTS_TOTAL .with_label_values(&[UNKNOWN_METHOD_LABEL, "error"]) .inc_by(0); WORKSPACE_RPC_ERRORS_TOTAL - .with_label_values(&[UNKNOWN_METHOD_LABEL, "hub_error"]) + .with_label_values(&[UNKNOWN_METHOD_LABEL, "unknown_method"]) .inc_by(0); let _ = WORKSPACE_RPC_DURATION_SECONDS.with_label_values(&[UNKNOWN_METHOD_LABEL]); } @@ -930,9 +926,7 @@ impl WorkspaceRpcHandler { } _ => { tracing::warn!(method, "unknown workspace rpc method"); - Err(WorkspaceError::HubError(format!( - "{UNKNOWN_METHOD_ERR_PREFIX} {method}" - ))) + Err(WorkspaceError::UnknownMethod(method.to_owned())) } } } @@ -989,10 +983,7 @@ impl ToolServerHandler for WorkspaceRpcHandler { bound_session.as_deref().map(|s| s.0.as_str()), ) .await; - let is_unknown_method = matches!( - &result, - Err(WorkspaceError::HubError(msg)) if msg.starts_with(UNKNOWN_METHOD_ERR_PREFIX) - ); + let is_unknown_method = matches!(&result, Err(WorkspaceError::UnknownMethod(_))); let method_label = if is_unknown_method { UNKNOWN_METHOD_LABEL } else { @@ -1273,15 +1264,18 @@ mod tests { assert_eq!(reply, turn_hook::HookReply::default()); } #[tokio::test] - async fn dispatch_unknown_method_returns_hub_error() { + async fn dispatch_unknown_method_returns_unknown_method_error() { let handle = make_handle(); let handler = WorkspaceRpcHandler::new(handle); let result = handler .dispatch("workspace.nonexistent", Value::Null, None) .await; - assert!( - matches!(result, Err(WorkspaceError::HubError(msg)) if msg.contains("unknown workspace method")) - ); + match result { + Err(WorkspaceError::UnknownMethod(method)) => { + assert_eq!(method, "workspace.nonexistent"); + } + other => panic!("expected UnknownMethod, got {other:?}"), + } } /// A hub evict runs the two-phase drain then settles into terminal /// ShuttingDown (not a lingering Draining) for an evicted workspace. @@ -2360,7 +2354,7 @@ mod tests { .with_label_values(&[UNKNOWN_METHOD_LABEL, "error"]) .get(); let kind_before = WORKSPACE_RPC_ERRORS_TOTAL - .with_label_values(&[UNKNOWN_METHOD_LABEL, "hub_error"]) + .with_label_values(&[UNKNOWN_METHOD_LABEL, "unknown_method"]) .get(); let mut stream = handler .handle_call( @@ -2378,7 +2372,7 @@ mod tests { ); assert!( WORKSPACE_RPC_ERRORS_TOTAL - .with_label_values(&[UNKNOWN_METHOD_LABEL, "hub_error"]) + .with_label_values(&[UNKNOWN_METHOD_LABEL, "unknown_method"]) .get() > kind_before, "a failed dispatch must also record its error_kind on the errors counter" diff --git a/crates/codegen/xai-grok-workspace/src/lib.rs b/crates/codegen/xai-grok-workspace/src/lib.rs index 230a5ba..34c109a 100644 --- a/crates/codegen/xai-grok-workspace/src/lib.rs +++ b/crates/codegen/xai-grok-workspace/src/lib.rs @@ -183,7 +183,7 @@ mod init_metrics_tests { )); assert!(has( "grok_workspace_rpc_errors_total", - &[("method", "unknown"), ("error_kind", "hub_error")] + &[("method", "unknown"), ("error_kind", "unknown_method")] )); for stage in [ "startup_recovery", diff --git a/crates/codegen/xai-grok-workspace/src/permission/manager.rs b/crates/codegen/xai-grok-workspace/src/permission/manager.rs index 89dfd48..854d7d0 100644 --- a/crates/codegen/xai-grok-workspace/src/permission/manager.rs +++ b/crates/codegen/xai-grok-workspace/src/permission/manager.rs @@ -5123,7 +5123,11 @@ mod tests { let local = tokio::task::LocalSet::new(); local .run_until(async { - for path in ["/etc/hosts", "/home/user/.grok/hooks/evil.json"] { + for path in [ + "/etc/hosts", + "/home/user/.grok/hooks/evil.json", + "/home/user/.grok/sandbox.toml", + ] { let mut auto = crate::permission::types::PermissionConfig::new(vec![]); auto.prompt_policy = PromptPolicy::Auto; let allow = diff --git a/crates/codegen/xai-grok-workspace/src/permission/shell_access.rs b/crates/codegen/xai-grok-workspace/src/permission/shell_access.rs index 45799d1..4be5ff0 100644 --- a/crates/codegen/xai-grok-workspace/src/permission/shell_access.rs +++ b/crates/codegen/xai-grok-workspace/src/permission/shell_access.rs @@ -328,6 +328,7 @@ pub enum ProtectedEditReason { StartupFile, Etc, GrokConfig, + GrokSandbox, ClaudeSettings, CursorHooks, /// Fail-closed / unclassified sensitive path; no user copy yet. @@ -343,6 +344,7 @@ impl ProtectedEditReason { Self::StartupFile => "startup_file", Self::Etc => "etc", Self::GrokConfig => "grok_config", + Self::GrokSandbox => "grok_sandbox", Self::ClaudeSettings => "claude_settings", Self::CursorHooks => "cursor_hooks", Self::Sensitive => "sensitive", @@ -369,6 +371,9 @@ impl ProtectedEditReason { Self::GrokConfig => Some( "Note: This edit contains changes to Grok config, which can alter permissions, tools, and other behavior in later sessions.", ), + Self::GrokSandbox => Some( + "Note: This edit contains changes to the Grok sandbox config, which can loosen filesystem and network restrictions on commands.", + ), Self::ClaudeSettings => Some( "Note: This edit contains changes to Claude-compatible settings, which can install hooks or change permission mode without a separate execution approval.", ), @@ -471,8 +476,8 @@ fn protected_edit_reason(path: &Path) -> Option { if STARTUP_FILES.contains(&file) { return Some(ProtectedEditReason::StartupFile); } - if string_components.ends_with(&[".grok", "config.toml"]) { - return Some(ProtectedEditReason::GrokConfig); + if let Some(reason) = protected_grok_config_file(path, &string_components) { + return Some(reason); } if path == Path::new("/etc") || path.starts_with(Path::new("/etc")) { return Some(ProtectedEditReason::Etc); @@ -480,6 +485,52 @@ fn protected_edit_reason(path: &Path) -> Option { None } +/// Grok config files that alter permissions (`config.toml`, the +/// `managed_config.toml` defaults tier, the user `requirements.toml` layer) or +/// sandbox restrictions (`sandbox.toml`) in the running and later sessions; a +/// silent edit would let the agent loosen its own guardrails. Matched directly +/// inside any `.grok` dir (user-global default and workspace overlays) and +/// directly under a custom `$GROK_HOME`, which the component match cannot see. +fn protected_grok_config_file(path: &Path, components: &[&str]) -> Option { + protected_grok_config_file_with_home( + path, + components, + xai_grok_config::user_grok_home().as_deref(), + ) +} + +fn protected_grok_config_file_with_home( + path: &Path, + components: &[&str], + user_grok_home: Option<&Path>, +) -> Option { + let reason = match components.last().copied() { + Some( + xai_grok_config::USER_CONFIG_FILENAME + | xai_grok_config::MANAGED_CONFIG_FILENAME + | xai_grok_config::REQUIREMENTS_FILENAME, + ) => ProtectedEditReason::GrokConfig, + Some("sandbox.toml") => ProtectedEditReason::GrokSandbox, + _ => return None, + }; + let in_dot_grok = components.len() >= 2 && components[components.len() - 2] == ".grok"; + let in_grok_home = || grok_home_matches(user_grok_home, |home| path.parent() == Some(home)); + (in_dot_grok || in_grok_home()).then_some(reason) +} + +/// True when `pred` holds for the user grok home in either its lexical or +/// physically-resolved form. Both forms are checked because callers hold a +/// lexical and a resolved candidate path, and the home itself may sit behind a +/// symlink. The comparison is byte-exact (no case folding), like every other +/// resolved-path check in this module. +fn grok_home_matches(home: Option<&Path>, pred: impl Fn(&Path) -> bool) -> bool { + home.is_some_and(|home| { + let lexical = xai_grok_paths::normalize_lexically(home); + pred(&lexical) + || resolve_following_symlinks(&lexical, 0).is_some_and(|resolved| pred(&resolved)) + }) +} + fn path_is_under_user_grok_hook_root(path: &Path, grok_home: &Path) -> bool { path.starts_with(grok_home.join("hooks")) || path == grok_home.join("hooks-paths") } @@ -487,12 +538,8 @@ fn path_is_under_user_grok_hook_root(path: &Path, grok_home: &Path) -> bool { fn protected_grok_hook_root(path: &Path, components: &[&str]) -> bool { components.windows(2).any(|pair| pair == [".grok", "hooks"]) || components.ends_with(&[".grok", "hooks-paths"]) - || xai_grok_config::user_grok_home().is_some_and(|grok_home| { - let lexical_home = xai_grok_paths::normalize_lexically(&grok_home); - path_is_under_user_grok_hook_root(path, &lexical_home) - || resolve_following_symlinks(&lexical_home, 0).is_some_and(|resolved_home| { - path_is_under_user_grok_hook_root(path, &resolved_home) - }) + || grok_home_matches(xai_grok_config::user_grok_home().as_deref(), |home| { + path_is_under_user_grok_hook_root(path, home) }) } @@ -1369,6 +1416,8 @@ mod tests { "/etc", "/etc/grok-test", "/work/subdir/../.git/hooks/pre-commit", + "/home/user/.grok/sandbox.toml", + "/work/project/.grok/sandbox.toml", ] { assert!( edit_target_protection(Path::new(path)).is_some(), @@ -1378,6 +1427,9 @@ mod tests { for path in [ "/work/src/main.rs", "/work/project/.grok/config.toml/backup", + "/work/project/sandbox.toml", + "/work/project/requirements.toml", + "/work/project/managed_config.toml", ] { assert!( edit_target_protection(Path::new(path)).is_none(), @@ -1428,6 +1480,22 @@ mod tests { "/home/user/.grok/config.toml", ProtectedEditReason::GrokConfig, ), + ( + "/home/user/.grok/sandbox.toml", + ProtectedEditReason::GrokSandbox, + ), + ( + "/work/project/.grok/sandbox.toml", + ProtectedEditReason::GrokSandbox, + ), + ( + "/home/user/.grok/managed_config.toml", + ProtectedEditReason::GrokConfig, + ), + ( + "/home/user/.grok/requirements.toml", + ProtectedEditReason::GrokConfig, + ), ( "/home/user/.claude/settings.json", ProtectedEditReason::ClaudeSettings, @@ -1551,6 +1619,85 @@ mod tests { } } + /// A custom `$GROK_HOME` has no `.grok` path component, so the live + /// `config.toml` / `sandbox.toml` must be caught by the home-prefix branch. + #[test] + fn grok_config_files_under_custom_grok_home_are_protected() { + let home = tempfile::tempdir().unwrap(); + let home_path = home.path(); + for (file, reason) in [ + ("config.toml", ProtectedEditReason::GrokConfig), + ("managed_config.toml", ProtectedEditReason::GrokConfig), + ("requirements.toml", ProtectedEditReason::GrokConfig), + ("sandbox.toml", ProtectedEditReason::GrokSandbox), + ] { + let path = home_path.join(file); + let components = [file]; + assert_eq!( + protected_grok_config_file_with_home(&path, &components, Some(home_path)), + Some(reason), + "{file} directly under $GROK_HOME must be protected" + ); + } + // Same file names elsewhere (or with no resolvable home) stay ordinary. + let elsewhere = home_path.join("sub").join("sandbox.toml"); + assert_eq!( + protected_grok_config_file_with_home( + &elsewhere, + &["sub", "sandbox.toml"], + Some(home_path) + ), + None + ); + assert_eq!( + protected_grok_config_file_with_home( + &home_path.join("sandbox.toml"), + &["sandbox.toml"], + None + ), + None + ); + } + + /// The resolved-symlink arm of the grok-home match must decide: `$GROK_HOME` + /// points at a symlink while the edit targets the physical home directory, + /// so the lexical parent-equality arm cannot fire. + #[test] + #[cfg(unix)] + fn grok_config_under_symlinked_grok_home_is_protected() { + use std::os::unix::fs::symlink; + let tmp = tempfile::tempdir().unwrap(); + let real_home = tmp.path().join("real-home"); + std::fs::create_dir(&real_home).unwrap(); + let link = tmp.path().join("home-link"); + symlink(&real_home, &link).unwrap(); + // tempdir paths can themselves contain symlinks (macOS /var -> /private/var); + // compare against the physical home the production resolver will produce. + let physical_home = resolve_following_symlinks(&real_home, 0).unwrap(); + assert_eq!( + protected_grok_config_file_with_home( + &physical_home.join("sandbox.toml"), + &["sandbox.toml"], + Some(&link) + ), + Some(ProtectedEditReason::GrokSandbox) + ); + } + + /// `protected_edit_reason` lowercases path components before matching, so + /// the canonical filename constants must stay lowercase or the const + /// patterns silently stop firing. + #[test] + fn protected_config_filename_constants_are_lowercase() { + for name in [ + xai_grok_config::USER_CONFIG_FILENAME, + xai_grok_config::MANAGED_CONFIG_FILENAME, + xai_grok_config::REQUIREMENTS_FILENAME, + ] { + assert_eq!(name, name.to_ascii_lowercase(), "{name}"); + } + } + #[test] fn resolved_root_alias_matches_physical_destination() { let resolved_root = resolve_following_symlinks(Path::new("/etc"), 0).unwrap(); diff --git a/crates/codegen/xai-grok-workspace/src/preview_supervisor.rs b/crates/codegen/xai-grok-workspace/src/preview_supervisor.rs index c5448b9..5cca68c 100644 --- a/crates/codegen/xai-grok-workspace/src/preview_supervisor.rs +++ b/crates/codegen/xai-grok-workspace/src/preview_supervisor.rs @@ -400,25 +400,39 @@ fn activity_url(control_port: u16) -> String { ) } +/// One scrape of the proxy's activity endpoint. `last_activity_ms` is the only +/// required field; the rest default to zero, so a workspace-server running +/// ahead of the proxy binary degrades to the old behaviour rather than failing. +#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, serde::Deserialize)] +struct ActivitySample { + last_activity_ms: u64, + #[serde(default)] + last_routed_ms: u64, + #[serde(default)] + ws_tunnels_open: u64, + #[serde(default)] + routed_requests_in_flight: u64, +} + /// Classified result of one scrape, so a missing proxy (quiet no-op) is never /// confused with a genuine error response or with real activity. #[derive(Debug, PartialEq, Eq)] enum ScrapeOutcome { - /// The proxy answered with a parseable activity stamp (epoch-ms). - Stamp(u64), + /// The proxy answered with a parseable activity sample. + Stamp(ActivitySample), /// The proxy isn't reachable (connection refused / not up yet): quiet no-op. Absent, /// The proxy answered but the response was unusable (error status / bad body). BadResponse, } -/// Parse `{ "last_activity_ms": }`; `None` for a malformed body, a missing -/// field, or a non-integer value. -fn parse_activity_body(body: &str) -> Option { - serde_json::from_str::(body) - .ok()? - .get("last_activity_ms")? - .as_u64() +/// Parse an activity body; `None` for a malformed body or a missing/non-integer +/// `last_activity_ms`. Unknown fields are ignored so the proxy can add more. +fn parse_activity_body(body: &str) -> Option { + let value: serde_json::Value = serde_json::from_str(body).ok()?; + // Required, and must be an integer. + value.get("last_activity_ms")?.as_u64()?; + serde_json::from_value(value).ok() } /// Classify a completed response by status + body (transport failures are @@ -428,7 +442,7 @@ fn classify_activity_response(status: u16, body: &str) -> ScrapeOutcome { return ScrapeOutcome::BadResponse; } match parse_activity_body(body) { - Some(ms) => ScrapeOutcome::Stamp(ms), + Some(sample) => ScrapeOutcome::Stamp(sample), None => ScrapeOutcome::BadResponse, } } @@ -457,6 +471,19 @@ async fn scrape_activity(client: &reqwest::Client, url: &str) -> ScrapeOutcome { } } +/// Clear the mirrored attached-client counters once we have been without +/// trustworthy data for `grace`. Starts the clock on the first bad scrape. +fn clear_attached_if_stale( + tracker: &ActivityTracker, + stale_since: &mut Option, + grace: Duration, +) { + let since = *stale_since.get_or_insert_with(Instant::now); + if since.elapsed() >= grace { + tracker.set_preview_attached(0, 0); + } +} + /// Poll the proxy's loopback activity endpoint until `shutdown` flips, feeding /// the tracker on each advance. Spawn after the `ActivityTracker` exists (post /// hub-connect), gated on preview being enabled. `control_port` is the proxy's @@ -507,22 +534,56 @@ async fn scrape_activity_loop( // `None` until the first successful scrape establishes a baseline. Baselining // (rather than starting at 0) avoids a spurious withhold when a workspace-server // restart meets a proxy whose stamp is already non-zero but stale. - let mut last_seen: Option = None; + let mut last_seen: Option = None; + // When we last had trustworthy attached-client data. The counters are the + // ONLY hold a WS-only client has — a tunnel writes no stamps — so one bad + // scrape must not clear them; but unlike the stamps they do not decay, so a + // sustained loss must, or a proxy that dies mid-tunnel holds the sandbox to + // the TTL. `Absent` and `BadResponse` both count as loss: one cannot reach + // the proxy, the other cannot understand it. The threshold is the stamp + // window, so both hold mechanisms expire on the same clock. + let attached_grace = Duration::from_millis(tracker.preview_activity_window_ms()); + let mut attached_stale_since: Option = None; loop { if sleep_or_shutdown(interval, &mut shutdown).await { return; } match scrape_activity(&client, &url).await { ScrapeOutcome::Stamp(current) => { - if last_seen.is_some_and(|prev| preview_activity_advanced(prev, current)) { - tracker.note_preview_activity(); + if let Some(prev) = last_seen { + // The routed stamp is a subset of the generic one, so check + // it first and attribute only the remainder to a poll — + // otherwise one routed request would report as both. + if preview_activity_advanced(prev.last_routed_ms, current.last_routed_ms) { + tracker.note_preview_routed_activity(); + } else if preview_activity_advanced( + prev.last_activity_ms, + current.last_activity_ms, + ) { + tracker.note_preview_status_activity(); + } } + // Absolute counters, republished every tick: a mirror of the + // proxy's state, not an edge. + tracker.set_preview_attached( + current.ws_tunnels_open, + current.routed_requests_in_flight, + ); last_seen = Some(current); + attached_stale_since = None; } - // Proxy absent (preview disabled / starting / restarting): no-op. - ScrapeOutcome::Absent => {} + // Proxy absent (preview disabled / starting / restarting): leave + // the stamps, and age the attached counters out. See above. + ScrapeOutcome::Absent => { + clear_attached_if_stale(&tracker, &mut attached_stale_since, attached_grace); + } + // Answering, but unusably. Same staleness clock: an error status or + // an unparseable body tells us nothing about attached clients, and + // leaving them untouched would let a persistently broken proxy hold + // the withhold open forever. ScrapeOutcome::BadResponse => { tracing::debug!(%url, "preview-activity scrape returned an unusable response"); + clear_attached_if_stale(&tracker, &mut attached_stale_since, attached_grace); } } } @@ -886,15 +947,23 @@ mod tests { .expect("a pre-flipped shutdown must return without scraping"); } + /// What an older proxy binary reports. + fn stamp_only(last_activity_ms: u64) -> ActivitySample { + ActivitySample { + last_activity_ms, + ..Default::default() + } + } + #[test] fn parse_activity_body_reads_stamp_and_rejects_bad_shapes() { assert_eq!( parse_activity_body(r#"{"last_activity_ms":1234}"#), - Some(1234) + Some(stamp_only(1234)) ); assert_eq!( parse_activity_body(r#"{"last_activity_ms":0,"extra":true}"#), - Some(0) + Some(stamp_only(0)) ); assert_eq!(parse_activity_body(r#"{"other":1}"#), None); assert_eq!(parse_activity_body(r#"{"last_activity_ms":"7"}"#), None); @@ -903,11 +972,130 @@ mod tests { assert_eq!(parse_activity_body(""), None); } + /// The binaries are version-pinned per session, but a restore can repin, so + /// this skew is real rather than theoretical. + #[test] + fn parse_activity_body_tolerates_a_proxy_without_the_new_fields() { + let old = parse_activity_body(r#"{"last_activity_ms":5,"status_holds_in_use":2}"#) + .expect("an old proxy body must still parse"); + assert_eq!(old.last_routed_ms, 0); + assert_eq!(old.ws_tunnels_open, 0); + assert_eq!( + old.routed_requests_in_flight, 0, + "absent fields read as 'nothing attached', never as attached" + ); + } + + /// A WS-only client has no stamps, so the attached counters are its only + /// hold. One unreachable scrape — a proxy restart, a loopback hiccup — must + /// not drop it and publish idle; a sustained absence still must. + #[tokio::test] + async fn one_absent_scrape_does_not_drop_an_attached_client() { + let tracker = Arc::new(ActivityTracker::new().with_preview_activity_window_ms(10_000)); + tracker.set_preview_attached(1, 0); + assert!(tracker.snapshot().idle_since_ms.is_none()); + + // Nothing is listening on this port, so every scrape classifies Absent. + let port = reserved_closed_port().await; + let (tx, rx) = watch::channel(false); + let loop_handle = tokio::spawn(scrape_activity_loop( + port, + Arc::clone(&tracker), + Duration::from_millis(10), + rx, + )); + + // Many consecutive absences, all well inside the 10s grace. + tokio::time::sleep(Duration::from_millis(300)).await; + assert!( + tracker.snapshot().idle_since_ms.is_none(), + "repeated unreachable scrapes inside the grace must not drop the hold" + ); + assert_eq!(tracker.preview_ws_tunnels_open(), 1); + + let _ = tx.send(true); + let _ = loop_handle.await; + } + + /// A proxy that answers unusably tells us nothing about attached clients + /// either, and the mirrored counters do not decay on their own — so an + /// endless run of error statuses must not pin `PreviewAttached` forever. + #[tokio::test] + async fn a_persistently_broken_proxy_does_not_pin_attached_forever() { + let tracker = Arc::new(ActivityTracker::new().with_preview_activity_window_ms(50)); + tracker.set_preview_attached(1, 0); + + // Answers every time, always with a 500 => BadResponse, never Absent. + let port = serve_canned("HTTP/1.1 500 Internal Server Error", "boom", true).await; + let (tx, rx) = watch::channel(false); + let loop_handle = tokio::spawn(scrape_activity_loop( + port, + Arc::clone(&tracker), + Duration::from_millis(10), + rx, + )); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + while tracker.preview_ws_tunnels_open() != 0 { + assert!( + tokio::time::Instant::now() < deadline, + "a sustained run of unusable responses must age the hold out" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } + + let _ = tx.send(true); + let _ = loop_handle.await; + } + + /// The other half: a proxy that stays gone must eventually release, or a + /// tunnel that died with it would hold the sandbox to the TTL. + #[tokio::test] + async fn a_sustained_absence_clears_the_attached_client() { + let tracker = Arc::new(ActivityTracker::new().with_preview_activity_window_ms(50)); + tracker.set_preview_attached(1, 0); + + let port = reserved_closed_port().await; + let (tx, rx) = watch::channel(false); + let loop_handle = tokio::spawn(scrape_activity_loop( + port, + Arc::clone(&tracker), + Duration::from_millis(10), + rx, + )); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + while tracker.preview_ws_tunnels_open() != 0 { + assert!( + tokio::time::Instant::now() < deadline, + "an absence past the window must clear the attached counters" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } + + let _ = tx.send(true); + let _ = loop_handle.await; + } + + #[test] + fn parse_activity_body_reads_the_attached_client_fields() { + let sample = parse_activity_body( + r#"{"last_activity_ms":9,"status_holds_in_use":1, + "held_status_aborts_quieted":0,"ws_tunnels_open":3, + "routed_requests_in_flight":2,"last_routed_ms":7}"#, + ) + .expect("parse"); + assert_eq!(sample.last_activity_ms, 9); + assert_eq!(sample.last_routed_ms, 7); + assert_eq!(sample.ws_tunnels_open, 3); + assert_eq!(sample.routed_requests_in_flight, 2); + } + #[test] fn classify_activity_response_distinguishes_stamp_from_bad() { assert_eq!( classify_activity_response(200, r#"{"last_activity_ms":42}"#), - ScrapeOutcome::Stamp(42) + ScrapeOutcome::Stamp(stamp_only(42)) ); assert_eq!( classify_activity_response(200, "garbage"), @@ -1022,7 +1210,7 @@ mod tests { let port = serve_canned("HTTP/1.1 200 OK", r#"{"last_activity_ms":9876}"#, false).await; assert_eq!( scrape_activity(&scrape_client(), &activity_url(port)).await, - ScrapeOutcome::Stamp(9876) + ScrapeOutcome::Stamp(stamp_only(9876)) ); } diff --git a/crates/codegen/xai-grok-workspace/src/rpc_envelope.rs b/crates/codegen/xai-grok-workspace/src/rpc_envelope.rs index 96f2716..e610286 100644 --- a/crates/codegen/xai-grok-workspace/src/rpc_envelope.rs +++ b/crates/codegen/xai-grok-workspace/src/rpc_envelope.rs @@ -30,6 +30,8 @@ pub fn error_code(err: &WorkspaceError) -> &'static str { WorkspaceError::InvalidHunkAction(_) => "invalid_hunk_action", WorkspaceError::HunkActionFailed(_) => "hunk_action_failed", WorkspaceError::HubError(_) => "hub_error", + WorkspaceError::UnknownMethod(_) => "unknown_method", + WorkspaceError::ExportArchiveLimitExceeded(_) => "export_archive_limit_exceeded", WorkspaceError::ExportGithub { kind, .. } => kind.wire_code(), WorkspaceError::ShuttingDown => "shutting_down", WorkspaceError::ToolsetExternallyOwned(_) => "toolset_externally_owned", @@ -80,6 +82,8 @@ pub fn rpc_error_to_workspace(err: RpcError) -> WorkspaceError { "invalid_hunk_action" => WorkspaceError::InvalidHunkAction(err.message), "hunk_action_failed" => WorkspaceError::HunkActionFailed(err.message), "hub_error" => WorkspaceError::HubError(err.message), + "unknown_method" => WorkspaceError::UnknownMethod(err.message), + "export_archive_limit_exceeded" => WorkspaceError::ExportArchiveLimitExceeded(err.message), "shutting_down" => WorkspaceError::ShuttingDown, "toolset_externally_owned" => WorkspaceError::ToolsetExternallyOwned(err.message), unknown => { @@ -115,6 +119,8 @@ mod tests { WorkspaceError::InvalidHunkAction("h".into()), WorkspaceError::HunkActionFailed("h".into()), WorkspaceError::HubError("hub".into()), + WorkspaceError::UnknownMethod("workspace.bogus".into()), + WorkspaceError::ExportArchiveLimitExceeded("too big".into()), WorkspaceError::ShuttingDown, WorkspaceError::ToolsetExternallyOwned("s".into()), ]; diff --git a/crates/common/xai-computer-hub-sdk/src/connection.rs b/crates/common/xai-computer-hub-sdk/src/connection.rs index a6d311a..179f6c4 100644 --- a/crates/common/xai-computer-hub-sdk/src/connection.rs +++ b/crates/common/xai-computer-hub-sdk/src/connection.rs @@ -50,12 +50,21 @@ use tracing::{info, warn}; use url::Url; use xai_tool_protocol::{ ConnectionId, ConnectionKind, JsonRpcId, JsonRpcRequest, JsonRpcResponse, JsonRpcVersion, - Method, PongFrame, ResponseOutcome, SessionId, + Method, PingFrame, PongFrame, ResponseOutcome, SessionId, }; /// Outbound mpsc bound. Picked to match the server's per-actor outbound /// buffer so a single-process roundtrip never dead-blocks on sender /// capacity. const OUTBOUND_BUFFER: usize = 256; +/// Writer control channel depth. Must fit a liveness `Close` + `Pause` +/// while still leaving room for `Resume` if the writer is mid-`sink.send`. +const WRITER_CTL_CAPACITY: usize = 4; +/// App-pong / priority outbound depth. Small and independent of the data +/// outbound buffer so heartbeats are not shed by tool-call backpressure. +const PRIORITY_OUTBOUND_CAPACITY: usize = 16; +/// Bound for the best-effort WS Close on a liveness kill. A silently dead +/// peer with a full TCP send buffer must not block Pause/Resume forever. +const WRITER_CLOSE_SEND_TIMEOUT: Duration = Duration::from_secs(2); /// Backoff schedule (in ms) for reconnect attempts. The last value is /// reused for any further attempts so the cap is `10s`. const RECONNECT_BACKOFF_MS: &[u64] = &[100, 200, 500, 1_000, 2_000, 5_000, 10_000]; @@ -86,9 +95,9 @@ struct HealthState { struct HealthSnapshot { last_inbound: Instant, /// Monotonic time elapsed since the last probe window rolled (the most - /// recent inbound frame or 5s clock probe) — NOT since connection start. - /// Healthy traffic keeps this small (<= ~5s); the meaningful freeze - /// signal in this snapshot is `clock_jump_ms`. + /// recent RTT proof — WS/app pong — or 5s clock probe) — NOT since + /// connection start. Healthy traffic keeps this small (<= ~5s); the + /// meaningful freeze signal in this snapshot is `clock_jump_ms`. since_last_probe_monotonic_ms: u64, /// Wall-clock time elapsed over the same probe window as /// `since_last_probe_monotonic_ms`. @@ -129,6 +138,9 @@ impl ConnHealth { state.mono_ref = Instant::now(); state.wall_ref = SystemTime::now(); } + /// Record RTT proof (WS/app pong). Hub→client pings/data must not call + /// this — they are one-way and would zero `detect_ms` / `silent_gap_ms` + /// during a mute that still expires the liveness deadline. fn record_inbound(&self) { let mut state = self.state.lock(); Self::roll(&mut state); @@ -171,8 +183,8 @@ enum DisconnectCause { ReadError(String), WriteError(String), Forced, - /// No inbound frame arrived within the inbound-liveness deadline, so the - /// transport is silently dead (snapshot-restored VM, NAT/LB flow expiry). + /// No RTT proof (WS/app pong or non-ping data) within the inbound-liveness + /// deadline — return path is silently dead. LivenessDeadline, } impl DisconnectCause { @@ -294,22 +306,20 @@ fn resolve_ws_ping_interval(configured: Option) -> Duration { } } /// Resolve the inbound-liveness deadline, clamping an unset *or zero* value -/// to 2.5× the (already-resolved) keepalive ping cadence — 75s at the -/// default 30s ping. +/// to `min(4× ping, 120s)` — 120s at the default 30s ping, still under the +/// hub's ~150s idle timeout. /// -/// The default multiple is chosen for fleet-wide false-positive safety: a -/// healthy connection delivers at least one inbound frame per ping period -/// (the server must answer each WS `Ping` with a `Pong`, and any data frame -/// also counts), so 2.5× tolerates a fully lost/coalesced pong plus -/// scheduling jitter before declaring death. It still detects a silently -/// dead transport (snapshot-restored VM, NAT/LB flow expiry) within ~1–2 -/// keepalive cycles instead of TCP-retransmission timescales (15+ min). -/// Explicit overrides are honored verbatim; keep them comfortably above -/// the ping interval or a healthy-but-idle connection will be churned. +/// After RTT-only re-arm, hub app/WS pings no longer keep the timer alive. +/// 4× (capped) tolerates a few lost/coalesced pongs plus scheduling jitter +/// without racing hub 4408. Explicit overrides are honored verbatim; keep +/// them comfortably above the ping interval or a healthy-but-idle +/// connection will be churned. fn resolve_ws_liveness_deadline(configured: Option, ping_interval: Duration) -> Duration { match configured { Some(deadline) if !deadline.is_zero() => deadline, - _ => ping_interval.saturating_mul(5) / 2, + _ => ping_interval + .saturating_mul(4) + .min(Duration::from_secs(120)), } } /// Optional, default-preserving connection-tuning knobs carried from the @@ -322,10 +332,11 @@ pub struct ConnectionTuning { /// Override for the keepalive ping cadence. `None` (or zero) ⇒ /// [`DEFAULT_WS_PING_INTERVAL`]. pub ws_ping_interval: Option, - /// Override for the inbound-liveness deadline: with no inbound frame of - /// any kind for this long, the reader declares the socket dead and - /// reconnects. `None` (or zero) ⇒ 2.5× the effective ping cadence (see - /// [`resolve_ws_liveness_deadline`]). + /// Override for the inbound-liveness deadline: with no *round-trip* + /// proof (WS/app pong) for this long, the reader declares the socket + /// dead and reconnects. Hub app pings and one-way hub→client data do + /// not re-arm. `None` (or zero) ⇒ `min(4× ping, 120s)` + /// (see [`resolve_ws_liveness_deadline`]). pub ws_liveness_deadline: Option, /// Override for the reconnect backoff schedule. `None` (or empty) ⇒ /// the built-in [`RECONNECT_BACKOFF_MS`] table. Stored as @@ -366,11 +377,12 @@ pub type ReconnectCallback = Box /// Boxed disconnect callback, fired when the live socket drops (before a /// reconnect attempt) and on a terminal close. pub type DisconnectCallback = Box; -/// Boxed connect callback, fired once on the initial successful connect, -/// before the reader actor task spawns. It therefore strictly happens-before -/// any disconnect/reconnect callback, so a connect/disconnect pair can never -/// be observed out of order (e.g. a readiness marker resurrected after the -/// socket has already dropped). +/// Boxed connect callback, fired once on the initial successful connect +/// after the writer keepalive loop has entered (so `/ready` cannot race +/// the first ping) and before the reader actor task spawns. It therefore +/// strictly happens-before any disconnect/reconnect callback, so a +/// connect/disconnect pair can never be observed out of order (e.g. a +/// readiness marker resurrected after the socket has already dropped). pub type ConnectCallback = Box; /// A live (or reconnecting) connection to the server. /// @@ -412,7 +424,8 @@ pub struct ConnectionConfig { /// server sends a terminal close. pub on_disconnect: Option>, /// Optional connect callback, fired once on the initial successful connect - /// before the actor starts (so it happens-before any disconnect/reconnect). + /// after the writer task enters its loop (happens-before reader start). + /// The first keepalive may still be in flight or one scheduler quanta away. pub on_connect: Option>, /// Stable server identity sent in the hello frame. Only meaningful /// for [`ConnectionKind::ToolServer`] connections. @@ -562,9 +575,6 @@ impl HubConnection { connection_id = %ack.connection_id, "server connection established" ); - if let Some(cb) = &config.on_connect { - cb(); - } let early_notif_rx = parking_lot::Mutex::new(match config.kind { ConnectionKind::ToolServer => Some(demux.subscribe_notifications()), _ => None, @@ -597,16 +607,24 @@ impl HubConnection { writer_error: writer_error.clone(), }); let (writer_ctl_tx, writer_ctl_rx) = - mpsc::channel::>>(2); + mpsc::channel::>>(WRITER_CTL_CAPACITY); let (writer_stop_tx, writer_stop_rx) = mpsc::channel::<()>(1); + let (priority_tx, priority_rx) = mpsc::channel::(PRIORITY_OUTBOUND_CAPACITY); + let (writer_ready_tx, writer_ready_rx) = oneshot::channel(); let writer_handle = tokio::spawn(run_writer( sink, outbound_rx, + priority_rx, writer_ctl_rx, writer_stop_rx, - ws_ping_interval, + Some(ws_ping_interval), writer_error, + Some(writer_ready_tx), )); + let _ = writer_ready_rx.await; + if let Some(cb) = &config.on_connect { + cb(); + } let reader_inner = inner.clone(); tokio::spawn(run_reader_actor( reader_inner, @@ -616,6 +634,7 @@ impl HubConnection { writer_ctl_tx, writer_stop_tx, writer_handle, + priority_tx, config.url, ws_liveness_deadline, )); @@ -1001,25 +1020,37 @@ fn now_unix_millis() -> u64 { .unwrap_or_default() .as_millis() as u64 } -/// Decode an inbound text frame. Returns the serialized [`PongFrame`] -/// to send back when the frame is an app-level server `ping`; otherwise -/// routes the frame through the demux and returns `None`. -fn route_or_pong(inner: &HubConnectionInner, text: &str) -> Option { +enum InboundText { + AppPing { pong: Option }, + AppPong, + Data, + Unparseable, +} +fn classify_inbound_text(inner: &HubConnectionInner, text: &str) -> InboundText { match serde_json::from_str::(text) { - Ok(value) => { - if value.get("method").and_then(Value::as_str) == Some(Method::Ping.as_wire_str()) { - serde_json::to_string(&PongFrame::new(now_unix_millis())).ok() - } else { + Ok(value) => match value.get("method").and_then(Value::as_str) { + Some(m) if m == Method::Ping.as_wire_str() => InboundText::AppPing { + pong: serde_json::to_string(&PongFrame::new(now_unix_millis())).ok(), + }, + Some(m) if m == Method::Pong.as_wire_str() => InboundText::AppPong, + _ => { let _ = inner.demux.route(value); - None + InboundText::Data } - } + }, Err(e) => { warn!(?e, "discarding unparseable inbound text frame"); - None + InboundText::Unparseable } } } +fn rearm_liveness(deadline: &mut std::pin::Pin<&mut tokio::time::Sleep>, liveness: Duration) { + let now = tokio::time::Instant::now(); + let rearm = now + .checked_add(liveness) + .unwrap_or_else(|| now + Duration::from_secs(86400 * 365 * 30)); + deadline.as_mut().reset(rearm); +} /// Map a websocket close frame's code to the connected-phase exit. Close /// codes 4100-4199 are terminal (the server intentionally ended the /// connection: eviction, session expiry, admin disconnect, rate limit). @@ -1053,90 +1084,303 @@ fn classify_stream_end(inner: &HubConnectionInner, read_error: Option) - /// The reader is the sole reconnect driver; it `Pause`s the writer the /// instant the socket is known dead so no buffered frame is dequeued /// onto the corpse, then `Resume`s it with the fresh sink once the -/// handshake completes. Carried on a cap-2 channel so a `Pause` is never -/// dropped. +/// handshake completes. Carried on [`WRITER_CTL_CAPACITY`] so a liveness +/// `Close`+`Pause` cannot crowd out `Resume`. enum WriterControl { /// Socket is dead; stop draining `outbound_rx` (frames stay buffered). Pause, /// Reconnected; install the fresh sink and resume draining. Resume(S), + /// Send a WS Close then stop draining (liveness kill / orderly drop). + Close { code: u16, reason: String }, +} +/// Outcome of racing a sink write against writer ctl/stop. +enum SendOrPreempt { + Sent(Result<(), String>), + Ctl(WriterControl), + Stop, +} +async fn send_or_preempt( + sink: &mut S, + msg: Message, + writer_ctl_rx: &mut mpsc::Receiver>, + writer_stop_rx: &mut mpsc::Receiver<()>, +) -> SendOrPreempt +where + S: futures::Sink + Unpin, + S::Error: std::fmt::Display, +{ + tokio::select! { + biased; + _ = writer_stop_rx.recv() => SendOrPreempt::Stop, + ctl = writer_ctl_rx.recv() => match ctl { + Some(ctl) => SendOrPreempt::Ctl(ctl), + None => SendOrPreempt::Stop, + }, + result = sink.send(msg) => SendOrPreempt::Sent(result.map_err(|e| e.to_string())), + } } /// Dedicated writer task: owns the sink, drains `outbound_rx`, and fires /// the keepalive ping (`ping_period`) — but only while `live`. Between a `Pause` and /// the matching `Resume` it parks on the control/stop channels only, so /// frames enqueued during the reconnect gap stay buffered in -/// `outbound_rx` and flush after `Resume` (no multi-frame loss; the -/// single in-flight frame whose `send` fails is the only loss, matching -/// the pre-split worst case). +/// `outbound_rx` and flush after `Resume`. /// -/// Generic over the sink so it can be unit-tested with an in-memory sink -/// without a live socket. +/// Data/ping writes are raced against ctl via [`send_or_preempt`] so a +/// half-open socket cannot strand Pause/Close/Resume behind `sink.send`. +/// Close is time-boxed against stop+timeout only — a queued Pause must +/// not abandon Close 1001. async fn run_writer( mut sink: S, mut outbound_rx: mpsc::Receiver, + mut priority_rx: mpsc::Receiver, mut writer_ctl_rx: mpsc::Receiver>, mut writer_stop_rx: mpsc::Receiver<()>, - ping_period: Duration, + ping_period: Option, write_error: WriteErrorSlot, + ready: Option>, ) where S: futures::Sink + Unpin, S::Error: std::fmt::Display, { - let mut ping_interval = tokio::time::interval(ping_period); - ping_interval.tick().await; + let keepalive = ping_period.is_some(); + let mut ping_interval = tokio::time::interval(ping_period.unwrap_or(Duration::from_secs(3600))); + ping_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + if !keepalive { + ping_interval.tick().await; + } let mut live = true; + let mut ready = ready; + let mut pending_app_ping = false; loop { - tokio::select! { - biased; - _ = writer_stop_rx.recv() => break, - ctl = writer_ctl_rx.recv() => match ctl { - Some(WriterControl::Pause) => live = false, - Some(WriterControl::Resume(new_sink)) => { - sink = new_sink; - live = true; - // Discard any error a late old-sink send left behind. The - // reader clears the slot before sending `Resume`, but an - // in-flight send on the dead socket (e.g. blocked on TCP - // retransmits since before `Pause`) can fail after that - // clear and re-fill the slot. This task is the only slot - // writer and processes messages sequentially, so by the - // time `Resume` is handled that old-sink send has - // finished — clearing here closes the race and stops a - // stale detail from mislabeling the NEXT disconnect as - // transport_write_error. - write_error.lock().take(); - // Restart the keepalive cadence from the reconnect instant: - // consume the immediate first tick so the next ping fires - // one period after Resume, not as a catch-up burst for ticks - // missed while paused. - ping_interval = tokio::time::interval(ping_period); - ping_interval.tick().await; + if let Some(tx) = ready.take() { + let _ = tx.send(()); + } + if live && pending_app_ping { + pending_app_ping = false; + let queued = match priority_rx.try_recv() { + Ok(text) => Some(text), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) => { + match outbound_rx.try_recv() { + Ok(text) => Some(text), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) => None, + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => { + break; + } + } } - // Reader gone (control sender dropped) → wind down. - None => break, - }, - _ = ping_interval.tick(), if live => { - if let Err(e) = sink.send(Message::Ping(Vec::new().into())).await { - // The reader detects the death (stream error, or liveness- - // deadline expiry once pings stop being answered) and - // drives the reconnect; we just stop draining onto the - // corpse. - *write_error.lock() = Some(format!("ping send failed: {e}")); - crate::metrics::writer_sink_send_error(); - live = false; - } - } - outbound = outbound_rx.recv(), if live => match outbound { - Some(text) => { - if let Err(e) = sink.send(Message::Text(text.into())).await { + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => break, + }; + let mut pending_ctl: Option> = None; + if let Some(text) = queued { + match send_or_preempt( + &mut sink, + Message::Text(text.into()), + &mut writer_ctl_rx, + &mut writer_stop_rx, + ) + .await + { + SendOrPreempt::Stop => break, + SendOrPreempt::Ctl(ctl) => pending_ctl = Some(ctl), + SendOrPreempt::Sent(Err(e)) => { *write_error.lock() = Some(format!("frame send failed: {e}")); crate::metrics::writer_sink_send_error(); live = false; } + SendOrPreempt::Sent(Ok(())) => {} } - // Last `outbound_tx` dropped → channel closed → wind down. + } + if live + && pending_ctl.is_none() + && let Ok(text) = serde_json::to_string(&PingFrame::new(now_unix_millis())) + { + match send_or_preempt( + &mut sink, + Message::Text(text.into()), + &mut writer_ctl_rx, + &mut writer_stop_rx, + ) + .await + { + SendOrPreempt::Stop => break, + SendOrPreempt::Ctl(ctl) => pending_ctl = Some(ctl), + SendOrPreempt::Sent(Err(e)) => { + *write_error.lock() = Some(format!("app ping send failed: {e}")); + crate::metrics::writer_sink_send_error(); + live = false; + } + SendOrPreempt::Sent(Ok(())) => {} + } + } + while let Some(ctl) = pending_ctl.take() { + match ctl { + WriterControl::Pause => { + live = false; + pending_app_ping = false; + while priority_rx.try_recv().is_ok() {} + } + WriterControl::Close { code, reason } => { + use tokio_tungstenite::tungstenite::protocol::CloseFrame; + use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; + live = false; + pending_app_ping = false; + while priority_rx.try_recv().is_ok() {} + let close_msg = Message::Close(Some(CloseFrame { + code: CloseCode::from(code), + reason: reason.into(), + })); + tokio::select! { + biased; + _ = writer_stop_rx.recv() => return, + _ = tokio::time::sleep(WRITER_CLOSE_SEND_TIMEOUT) => { + *write_error.lock() = + Some("close send timed out".to_owned()); + crate::metrics::writer_sink_send_error(); + } + result = sink.send(close_msg) => { + if let Err(e) = result { + *write_error.lock() = + Some(format!("close send failed: {e}")); + crate::metrics::writer_sink_send_error(); + } + } + } + } + WriterControl::Resume(new_sink) => { + sink = new_sink; + live = true; + pending_app_ping = false; + while priority_rx.try_recv().is_ok() {} + write_error.lock().take(); + ping_interval = + tokio::time::interval(ping_period.unwrap_or(Duration::from_secs(3600))); + ping_interval + .set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + if !keepalive { + ping_interval.tick().await; + } + } + } + } + continue; + } + let mut pending_ctl: Option> = tokio::select! { + biased; + _ = writer_stop_rx.recv() => break, + ctl = writer_ctl_rx.recv() => match ctl { + Some(ctl) => Some(ctl), None => break, }, + _ = ping_interval.tick(), if live && keepalive && !pending_app_ping => { + match send_or_preempt( + &mut sink, + Message::Ping(Vec::new().into()), + &mut writer_ctl_rx, + &mut writer_stop_rx, + ).await { + SendOrPreempt::Stop => break, + SendOrPreempt::Ctl(ctl) => Some(ctl), + SendOrPreempt::Sent(Err(e)) => { + *write_error.lock() = Some(format!("ping send failed: {e}")); + crate::metrics::writer_sink_send_error(); + live = false; + None + } + SendOrPreempt::Sent(Ok(())) => { + // App ping survives proxies that eat WS control frames. + pending_app_ping = true; + None + } + } + } + priority = priority_rx.recv(), if live => match priority { + Some(text) => match send_or_preempt( + &mut sink, + Message::Text(text.into()), + &mut writer_ctl_rx, + &mut writer_stop_rx, + ).await { + SendOrPreempt::Stop => break, + SendOrPreempt::Ctl(ctl) => Some(ctl), + SendOrPreempt::Sent(Err(e)) => { + *write_error.lock() = Some(format!("priority send failed: {e}")); + crate::metrics::writer_sink_send_error(); + live = false; + None + } + SendOrPreempt::Sent(Ok(())) => None, + }, + None => break, + }, + outbound = outbound_rx.recv(), if live => match outbound { + Some(text) => match send_or_preempt( + &mut sink, + Message::Text(text.into()), + &mut writer_ctl_rx, + &mut writer_stop_rx, + ).await { + SendOrPreempt::Stop => break, + SendOrPreempt::Ctl(ctl) => Some(ctl), + SendOrPreempt::Sent(Err(e)) => { + *write_error.lock() = Some(format!("frame send failed: {e}")); + crate::metrics::writer_sink_send_error(); + live = false; + None + } + SendOrPreempt::Sent(Ok(())) => None, + }, + None => break, + }, + }; + while let Some(ctl) = pending_ctl.take() { + match ctl { + WriterControl::Pause => { + live = false; + pending_app_ping = false; + while priority_rx.try_recv().is_ok() {} + } + WriterControl::Close { code, reason } => { + use tokio_tungstenite::tungstenite::protocol::CloseFrame; + use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; + live = false; + pending_app_ping = false; + while priority_rx.try_recv().is_ok() {} + let close_msg = Message::Close(Some(CloseFrame { + code: CloseCode::from(code), + reason: reason.into(), + })); + tokio::select! { + biased; + _ = writer_stop_rx.recv() => return, + _ = tokio::time::sleep(WRITER_CLOSE_SEND_TIMEOUT) => { + *write_error.lock() = + Some("close send timed out".to_owned()); + crate::metrics::writer_sink_send_error(); + } + result = sink.send(close_msg) => { + if let Err(e) = result { + *write_error.lock() = + Some(format!("close send failed: {e}")); + crate::metrics::writer_sink_send_error(); + } + } + } + } + WriterControl::Resume(new_sink) => { + sink = new_sink; + live = true; + pending_app_ping = false; + while priority_rx.try_recv().is_ok() {} + write_error.lock().take(); + ping_interval = + tokio::time::interval(ping_period.unwrap_or(Duration::from_secs(3600))); + ping_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + if !keepalive { + ping_interval.tick().await; + } + } + } } } } @@ -1157,6 +1401,7 @@ async fn run_reader_actor( writer_ctl_tx: mpsc::Sender>>, writer_stop_tx: mpsc::Sender<()>, writer_handle: tokio::task::JoinHandle<()>, + priority_tx: mpsc::Sender, url: Url, liveness_deadline: Duration, ) { @@ -1169,6 +1414,7 @@ async fn run_reader_actor( &mut stop_rx, &mut reconnect_rx, liveness_deadline, + &priority_tx, ) .await { @@ -1212,6 +1458,17 @@ async fn run_reader_actor( "server connection lost; scheduling reconnect" ); fire_on_disconnect(inner.as_ref()); + if matches!(outage.cause, DisconnectCause::LivenessDeadline) + && writer_ctl_tx + .send(WriterControl::Close { + code: 1001, + reason: "liveness_deadline".to_owned(), + }) + .await + .is_err() + { + break; + } if writer_ctl_tx.send(WriterControl::Pause).await.is_err() { break; } @@ -1322,11 +1579,11 @@ fn drain_reconnect_signals(reconnect_rx: &mut mpsc::Receiver<()>) { /// half but never writes (app-level pongs route through `outbound_tx`; WS /// pings are auto-answered by tungstenite on poll). /// -/// Enforces the inbound-liveness deadline: no inbound frame of any kind for -/// the deadline window (default 2.5× the ping cadence, see -/// [`resolve_ws_liveness_deadline`]) means the transport is silently dead -/// (snapshot-restored VM, NAT/LB flow expiry), so exit via -/// [`ConnectedExit::SocketClosed`] onto the normal reconnect path. The +/// Enforces the inbound-liveness deadline: no *round-trip* proof (WS/app +/// pong) for the deadline window (default 4× the ping +/// cadence, see [`resolve_ws_liveness_deadline`]) means the return path is +/// silently dead, so exit via [`ConnectedExit::SocketClosed`] onto the +/// normal reconnect path. Hub app/WS pings alone do not re-arm. The /// deadline runs only in this phase and re-arms on every (re)entry. /// /// Generic over the stream for in-memory unit tests, mirroring @@ -1337,11 +1594,13 @@ async fn run_reader_phase( stop_rx: &mut mpsc::Receiver<()>, reconnect_rx: &mut mpsc::Receiver<()>, liveness_deadline: Duration, + pong_tx: &mpsc::Sender, ) -> ConnectedExit where S: Stream> + Unpin, { let mut clock_probe = tokio::time::interval(CLOCK_PROBE_INTERVAL); + clock_probe.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); clock_probe.tick().await; let deadline = sleep(liveness_deadline); tokio::pin!(deadline); @@ -1356,37 +1615,37 @@ where // Before the deadline arm so a frame that raced the expiry // proves liveness and wins. msg = stream.next() => { - if matches!(msg, Some(Ok(ref m)) if !matches!(m, Message::Close(_))) { - inner.health.record_inbound(); - } match msg { Some(Ok(msg)) => { - // Any inbound frame (data or control) proves liveness, - // so re-arm the deadline. Saturate on overflow so a - // `Duration::MAX` "disable" override can't panic - // `Instant + Duration`. - let now = tokio::time::Instant::now(); - let rearm = now - .checked_add(liveness_deadline) - .unwrap_or_else(|| now + Duration::from_secs(86400 * 365 * 30)); - deadline.as_mut().reset(rearm); match msg { Message::Text(text) => { - if let Some(pong_text) = route_or_pong(inner, text.as_ref()) - && inner.outbound_tx.try_send(pong_text).is_err() - { - // App-level pong is JSON text; the reader no longer - // owns the sink, so route it through the writer. - // Best-effort (non-blocking) to keep the reader hot: - // a paused writer (dead socket) or a saturated buffer - // drops the heartbeat. Metered so the residual loss is - // observable/alertable rather than silent. - crate::metrics::heartbeat_pong_dropped(); + match classify_inbound_text(inner, text.as_ref()) { + InboundText::AppPing { pong } => { + // Hub→client ping is not RTT proof. Reply if we + // can; a dropped pong must not re-arm either. + if let Some(pong_text) = pong + && pong_tx.try_send(pong_text).is_err() + { + crate::metrics::heartbeat_pong_dropped(); + } + } + InboundText::AppPong => { + inner.health.record_inbound(); + rearm_liveness(&mut deadline, liveness_deadline); + } + InboundText::Data => { + // Hub→client data is one-way, not RTT proof. + } + InboundText::Unparseable => {} } } - // WS control pings get an automatic Pong queued + flushed - // by tungstenite on read; nothing to do here. - Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => {} + // WS Pong is RTT proof of our Ping. Inbound WS Ping is + // auto-answered by tungstenite and is hub→client only. + Message::Pong(_) => { + inner.health.record_inbound(); + rearm_liveness(&mut deadline, liveness_deadline); + } + Message::Ping(_) | Message::Frame(_) => {} Message::Binary(_) => { warn!("server sent binary frame; ignoring"); } @@ -1411,7 +1670,7 @@ where crate::metrics::liveness_deadline_expired(); warn!( ?liveness_deadline, - "no inbound frame within the liveness deadline; declaring the socket dead and reconnecting" + "no RTT proof (WS/app pong) within the liveness deadline; declaring the socket dead and reconnecting" ); return ConnectedExit::SocketClosed(DisconnectCause::LivenessDeadline); } @@ -1775,9 +2034,6 @@ mod tests { use std::pin::Pin; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::task::{Context, Poll}; - /// Default ping period for writer tests that don't exercise the - /// keepalive: long enough that no ping fires during the test. - const TEST_PING_NEVER: Duration = Duration::from_secs(3_600); /// In-memory [`futures::Sink`] for `run_writer` tests. Records the /// text payload of every `Message::Text` sent and counts every /// `Message::Ping` (keepalive). When the `fail` flag is set, `send` @@ -1832,6 +2088,15 @@ mod tests { Message::Ping(_) => { self.pings.fetch_add(1, Ordering::SeqCst); } + Message::Close(frame) => { + let (code, reason) = frame + .map(|f| (u16::from(f.code), f.reason.to_string())) + .unwrap_or((0, String::new())); + self.recorded + .lock() + .expect("recorded lock") + .push(format!("CLOSE:{code}:{reason}")); + } _ => {} } Ok(()) @@ -1853,6 +2118,15 @@ mod tests { fn idle_write_error_slot() -> WriteErrorSlot { Arc::new(parking_lot::Mutex::new(None)) } + fn outbound_data_frames(recorded: &std::sync::Mutex>) -> Vec { + recorded + .lock() + .expect("lock") + .iter() + .filter(|f| !(f.contains("\"method\":\"ping\"") && f.contains("ts_ms"))) + .cloned() + .collect() + } /// Poll `predicate` every 5ms up to ~2s. Keeps the writer-task tests /// off arbitrary fixed sleeps for the positive assertions. async fn wait_until bool>(predicate: F, label: &str) { @@ -1865,55 +2139,537 @@ mod tests { panic!("timed out waiting for: {label}"); } #[tokio::test] + async fn writer_sends_close_before_pause() { + let sink = RecordingSink::new(); + let recorded = sink.recorded(); + let (out_tx, out_rx) = mpsc::channel::(8); + let (ctl_tx, ctl_rx) = mpsc::channel::(4); + let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); + let writer = tokio::spawn(run_writer( + sink, + out_rx, + prio_rx, + ctl_rx, + stop_rx, + None, + idle_write_error_slot(), + None, + )); + ctl_tx + .send(WriterControl::Close { + code: 1001, + reason: "liveness_deadline".to_owned(), + }) + .await + .expect("close"); + ctl_tx.send(WriterControl::Pause).await.expect("pause"); + out_tx.send("buffered".to_owned()).await.expect("buffer"); + wait_until( + || { + recorded + .lock() + .expect("lock") + .iter() + .any(|f| f == "CLOSE:1001:liveness_deadline") + }, + "close frame written", + ) + .await; + assert!( + !recorded + .lock() + .expect("lock") + .iter() + .any(|f| f == "buffered"), + "buffered data must not flush on the old sink after Close+Pause" + ); + let fresh = RecordingSink::new(); + let fresh_recorded = fresh.recorded(); + ctl_tx + .send(WriterControl::Resume(fresh)) + .await + .expect("resume"); + wait_until( + || { + fresh_recorded + .lock() + .expect("lock") + .iter() + .any(|f| f == "buffered") + }, + "buffered data flushes on the fresh sink after Resume", + ) + .await; + drop(out_tx); + stop_tx.send(()).await.expect("stop"); + writer.await.expect("writer task joins"); + } + /// Sink whose first non-Close `poll_ready` stays pending until released. + /// Models a half-open peer with a full TCP send buffer. + struct BlockingSink { + recorded: Arc>>, + block: Arc, + waker: Arc>>, + } + impl Clone for BlockingSink { + fn clone(&self) -> Self { + Self { + recorded: self.recorded.clone(), + block: self.block.clone(), + waker: self.waker.clone(), + } + } + } + impl BlockingSink { + fn new() -> Self { + Self { + recorded: Arc::new(std::sync::Mutex::new(Vec::new())), + block: Arc::new(AtomicBool::new(true)), + waker: Arc::new(std::sync::Mutex::new(None)), + } + } + fn recorded(&self) -> Arc>> { + self.recorded.clone() + } + fn release(&self) { + self.block.store(false, Ordering::SeqCst); + if let Some(w) = self.waker.lock().expect("waker").take() { + w.wake(); + } + } + } + impl futures::Sink for BlockingSink { + type Error = std::io::Error; + fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.block.load(Ordering::SeqCst) { + *self.waker.lock().expect("waker") = Some(cx.waker().clone()); + Poll::Pending + } else { + Poll::Ready(Ok(())) + } + } + fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { + match item { + Message::Text(text) => self + .recorded + .lock() + .expect("lock") + .push(text.as_str().to_owned()), + Message::Close(frame) => { + let (code, reason) = frame + .map(|f| (u16::from(f.code), f.reason.to_string())) + .unwrap_or((0, String::new())); + self.recorded + .lock() + .expect("lock") + .push(format!("CLOSE:{code}:{reason}")); + } + _ => {} + } + Ok(()) + } + fn poll_flush( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + fn poll_close( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + } + #[tokio::test] + async fn writer_preempts_blocked_data_send_for_close() { + let sink = BlockingSink::new(); + let recorded = sink.recorded(); + let (out_tx, out_rx) = mpsc::channel::(8); + let (ctl_tx, ctl_rx) = mpsc::channel::>(4); + let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); + let writer = tokio::spawn(run_writer( + sink.clone(), + out_rx, + prio_rx, + ctl_rx, + stop_rx, + None, + idle_write_error_slot(), + None, + )); + out_tx.send("stuck".to_owned()).await.expect("data"); + tokio::time::sleep(Duration::from_millis(20)).await; + ctl_tx + .send(WriterControl::Close { + code: 1001, + reason: "liveness_deadline".to_owned(), + }) + .await + .expect("close"); + ctl_tx.send(WriterControl::Pause).await.expect("pause"); + tokio::time::sleep(Duration::from_millis(10)).await; + sink.release(); + wait_until( + || { + recorded + .lock() + .expect("lock") + .iter() + .any(|f| f == "CLOSE:1001:liveness_deadline") + }, + "close preempts blocked data write", + ) + .await; + assert!( + !recorded.lock().expect("lock").iter().any(|f| f == "stuck"), + "blocked data must not be written after Close preempt" + ); + drop(out_tx); + stop_tx.send(()).await.expect("stop"); + writer.await.expect("writer task joins"); + } + /// Accepts frames (records them) but never completes flush — models a + /// half-open TCP sndbuf so Close can be observed then time out. + #[tokio::test] + async fn writer_close_then_queued_pause_still_writes_close_1001() { + let sink = RecordingSink::new(); + let recorded = sink.recorded(); + let (out_tx, out_rx) = mpsc::channel::(4); + let (ctl_tx, ctl_rx) = mpsc::channel::(4); + let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); + let writer = tokio::spawn(run_writer( + sink, + out_rx, + prio_rx, + ctl_rx, + stop_rx, + None, + idle_write_error_slot(), + None, + )); + ctl_tx + .send(WriterControl::Close { + code: 1001, + reason: "liveness_deadline".to_owned(), + }) + .await + .expect("close"); + ctl_tx.send(WriterControl::Pause).await.expect("pause"); + wait_until( + || { + recorded + .lock() + .expect("lock") + .iter() + .any(|f| f.starts_with("CLOSE:1001:")) + }, + "Close 1001 recorded", + ) + .await; + let live = RecordingSink::new(); + let live_log = live.recorded(); + ctl_tx + .send(WriterControl::Resume(live)) + .await + .expect("resume"); + out_tx.send("after".to_owned()).await.expect("after"); + wait_until( + || outbound_data_frames(&live_log).iter().any(|f| f == "after"), + "Resume installs a live sink", + ) + .await; + stop_tx.send(()).await.expect("stop"); + writer.await.expect("join"); + } + #[tokio::test] + async fn writer_ctl_preempts_in_flight_blocking_ping_then_resumes() { + let sink = RecordingSink::new(); + let (out_tx, out_rx) = mpsc::channel::(4); + let (ctl_tx, ctl_rx) = mpsc::channel::(4); + let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); + let writer = tokio::spawn(run_writer( + sink, + out_rx, + prio_rx, + ctl_rx, + stop_rx, + Some(Duration::from_millis(1)), + idle_write_error_slot(), + None, + )); + ctl_tx + .send(WriterControl::Close { + code: 1001, + reason: "liveness_deadline".to_owned(), + }) + .await + .expect("close"); + ctl_tx.send(WriterControl::Pause).await.expect("pause"); + let live = RecordingSink::new(); + let live_log = live.recorded(); + ctl_tx + .send(WriterControl::Resume(live)) + .await + .expect("resume"); + out_tx.send("resumed".to_owned()).await.expect("send"); + wait_until( + || { + outbound_data_frames(&live_log) + .iter() + .any(|f| f == "resumed") + }, + "fresh sink accepts data after close+pause+resume", + ) + .await; + stop_tx.send(()).await.expect("stop"); + writer.await.expect("join"); + } + #[tokio::test] async fn writer_drains_outbound_while_live() { let sink = RecordingSink::new(); let recorded = sink.recorded(); let (out_tx, out_rx) = mpsc::channel::(8); let (_ctl_tx, ctl_rx) = mpsc::channel::(2); let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); let writer = tokio::spawn(run_writer( sink, out_rx, + prio_rx, ctl_rx, stop_rx, - TEST_PING_NEVER, + None, idle_write_error_slot(), + None, )); out_tx.send("a".to_owned()).await.expect("send a"); out_tx.send("b".to_owned()).await.expect("send b"); wait_until( - || recorded.lock().expect("lock").len() == 2, + || outbound_data_frames(&recorded).len() == 2, "two frames drained", ) .await; assert_eq!( - *recorded.lock().expect("lock"), + outbound_data_frames(&recorded), vec!["a".to_owned(), "b".to_owned()], "frames must be written to the live sink in order" ); stop_tx.send(()).await.expect("stop"); writer.await.expect("writer task joins"); } + #[tokio::test(start_paused = true)] + async fn writer_first_ping_is_immediate() { + let sink = RecordingSink::new(); + let pings = sink.pings(); + let recorded = sink.recorded(); + let (_out_tx, out_rx) = mpsc::channel::(4); + let (_ctl_tx, ctl_rx) = mpsc::channel::(2); + let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); + let writer = tokio::spawn(run_writer( + sink, + out_rx, + prio_rx, + ctl_rx, + stop_rx, + Some(Duration::from_secs(30)), + idle_write_error_slot(), + None, + )); + tokio::time::advance(Duration::from_millis(1)).await; + tokio::task::yield_now().await; + assert!( + pings.load(Ordering::SeqCst) >= 1, + "WS ping must fire immediately after writer spawn" + ); + assert!( + recorded + .lock() + .expect("lock") + .iter() + .any(|f| f.contains("\"method\":\"ping\"")), + "app ping must fire immediately after writer spawn" + ); + stop_tx.send(()).await.expect("stop"); + writer.await.expect("writer task joins"); + } + #[tokio::test(start_paused = true)] + async fn writer_first_ping_after_resume_is_immediate() { + let dead = RecordingSink::new(); + let (_out_tx, out_rx) = mpsc::channel::(4); + let (ctl_tx, ctl_rx) = mpsc::channel::(4); + let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); + let writer = tokio::spawn(run_writer( + dead, + out_rx, + prio_rx, + ctl_rx, + stop_rx, + Some(Duration::from_secs(30)), + idle_write_error_slot(), + None, + )); + ctl_tx.send(WriterControl::Pause).await.expect("pause"); + let fresh = RecordingSink::new(); + let pings = fresh.pings(); + let recorded = fresh.recorded(); + ctl_tx + .send(WriterControl::Resume(fresh)) + .await + .expect("resume"); + tokio::time::advance(Duration::from_millis(1)).await; + tokio::task::yield_now().await; + assert!( + pings.load(Ordering::SeqCst) >= 1, + "WS ping must fire immediately after Resume" + ); + assert!( + recorded + .lock() + .expect("lock") + .iter() + .any(|f| f.contains("\"method\":\"ping\"")), + "app ping must fire immediately after Resume" + ); + stop_tx.send(()).await.expect("stop"); + writer.await.expect("writer task joins"); + } + #[tokio::test] + async fn writer_priority_pong_bypasses_full_outbound() { + let sink = RecordingSink::new(); + let recorded = sink.recorded(); + let (out_tx, out_rx) = mpsc::channel::(1); + let (prio_tx, prio_rx) = mpsc::channel::(4); + let (_ctl_tx, ctl_rx) = mpsc::channel::(2); + let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + out_tx + .try_send("blocked".to_owned()) + .expect("fill outbound"); + prio_tx + .try_send(r#"{"method":"pong","ts_ms":1}"#.to_owned()) + .expect("priority send before spawn"); + let writer = tokio::spawn(run_writer( + sink, + out_rx, + prio_rx, + ctl_rx, + stop_rx, + None, + idle_write_error_slot(), + None, + )); + wait_until( + || { + outbound_data_frames(&recorded) + .first() + .is_some_and(|f| f.contains("\"method\":\"pong\"")) + }, + "priority pong is the first data frame", + ) + .await; + assert!( + outbound_data_frames(&recorded) + .iter() + .any(|f| f.contains("\"method\":\"pong\"")), + ); + assert!( + outbound_data_frames(&recorded) + .iter() + .position(|f| f.contains("\"method\":\"pong\"")) + < outbound_data_frames(&recorded) + .iter() + .position(|f| f == "blocked"), + ); + stop_tx.send(()).await.expect("stop"); + writer.await.expect("writer task joins"); + } + #[tokio::test] + async fn writer_pause_drops_stale_priority_pongs() { + let dead = RecordingSink::new(); + let (out_tx, out_rx) = mpsc::channel::(4); + let (prio_tx, prio_rx) = mpsc::channel::(4); + let (ctl_tx, ctl_rx) = mpsc::channel::(4); + let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + let writer = tokio::spawn(run_writer( + dead, + out_rx, + prio_rx, + ctl_rx, + stop_rx, + None, + idle_write_error_slot(), + None, + )); + ctl_tx.send(WriterControl::Pause).await.expect("pause"); + tokio::time::sleep(Duration::from_millis(20)).await; + prio_tx + .try_send(r#"{"method":"pong","ts_ms":1}"#.to_owned()) + .expect("stale pong"); + let fresh = RecordingSink::new(); + let recorded = fresh.recorded(); + ctl_tx + .send(WriterControl::Resume(fresh)) + .await + .expect("resume"); + tokio::time::sleep(Duration::from_millis(30)).await; + assert!( + !recorded + .lock() + .expect("lock") + .iter() + .any(|f| f.contains("pong")), + "stale priority pong must be drained on Pause/Resume" + ); + drop(out_tx); + stop_tx.send(()).await.expect("stop"); + writer.await.expect("join"); + } #[tokio::test] async fn writer_honors_custom_ping_interval() { let sink = RecordingSink::new(); let pings = sink.pings(); + let recorded = sink.recorded(); let (_out_tx, out_rx) = mpsc::channel::(4); let (_ctl_tx, ctl_rx) = mpsc::channel::(2); let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); let writer = tokio::spawn(run_writer( sink, out_rx, + prio_rx, ctl_rx, stop_rx, - Duration::from_millis(20), + Some(Duration::from_millis(20)), idle_write_error_slot(), + None, )); wait_until( || pings.load(Ordering::SeqCst) >= 3, "three keepalive pings at the configured cadence", ) .await; + wait_until( + || { + recorded.lock().expect("lock").iter().any(|text| { + serde_json::from_str::(text) + .ok() + .and_then(|v| { + v.get("method") + .and_then(serde_json::Value::as_str) + .map(|m| m == "ping") + }) + .unwrap_or(false) + }) + }, + "serialized app ping {\"method\":\"ping\",...} on the sink", + ) + .await; stop_tx.send(()).await.expect("stop"); writer.await.expect("writer task joins"); } @@ -1923,13 +2679,16 @@ mod tests { let (_out_tx, out_rx) = mpsc::channel::(4); let (ctl_tx, ctl_rx) = mpsc::channel::(2); let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); let writer = tokio::spawn(run_writer( dead, out_rx, + prio_rx, ctl_rx, stop_rx, - Duration::from_millis(20), + Some(Duration::from_millis(20)), idle_write_error_slot(), + None, )); ctl_tx.send(WriterControl::Pause).await.expect("pause"); let fresh = RecordingSink::new(); @@ -1953,13 +2712,16 @@ mod tests { let (out_tx, out_rx) = mpsc::channel::(16); let (ctl_tx, ctl_rx) = mpsc::channel::(2); let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); let writer = tokio::spawn(run_writer( dead, out_rx, + prio_rx, ctl_rx, stop_rx, - TEST_PING_NEVER, + None, idle_write_error_slot(), + None, )); ctl_tx.send(WriterControl::Pause).await.expect("pause"); tokio::time::sleep(Duration::from_millis(20)).await; @@ -1971,9 +2733,9 @@ mod tests { } tokio::time::sleep(Duration::from_millis(50)).await; assert!( - dead_log.lock().expect("lock").is_empty(), + outbound_data_frames(&dead_log).is_empty(), "paused writer must not drain onto the dead sink; got {:?}", - dead_log.lock().expect("lock") + outbound_data_frames(&dead_log) ); let fresh = RecordingSink::new(); let fresh_log = fresh.recorded(); @@ -1982,18 +2744,18 @@ mod tests { .await .expect("resume"); wait_until( - || fresh_log.lock().expect("lock").len() == 3, + || outbound_data_frames(&fresh_log).len() == 3, "buffered frames flush after resume", ) .await; assert_eq!( - *fresh_log.lock().expect("lock"), + outbound_data_frames(&fresh_log), vec!["g1".to_owned(), "g2".to_owned(), "g3".to_owned()], "all gap frames flush, in order, to the fresh sink" ); assert!( - dead_log.lock().expect("lock").is_empty(), - "no frame must ever reach the dead sink" + outbound_data_frames(&dead_log).is_empty(), + "no data frame must ever reach the dead sink" ); stop_tx.send(()).await.expect("stop"); writer.await.expect("writer task joins"); @@ -2007,17 +2769,20 @@ mod tests { let (ctl_tx, ctl_rx) = mpsc::channel::(2); let (stop_tx, stop_rx) = mpsc::channel::<()>(1); let write_error = idle_write_error_slot(); + let (_prio_tx, prio_rx) = mpsc::channel::(4); let writer = tokio::spawn(run_writer( failing, out_rx, + prio_rx, ctl_rx, stop_rx, - TEST_PING_NEVER, + None, write_error.clone(), + None, )); out_tx.send("ok".to_owned()).await.expect("send ok"); wait_until( - || failing_log.lock().expect("lock").len() == 1, + || outbound_data_frames(&failing_log).len() == 1, "first frame drained before failure", ) .await; @@ -2033,7 +2798,7 @@ mod tests { .expect("enqueue kept2"); tokio::time::sleep(Duration::from_millis(50)).await; assert_eq!( - *failing_log.lock().expect("lock"), + outbound_data_frames(&failing_log), vec!["ok".to_owned()], "only the pre-failure frame should have been recorded on the dead sink" ); @@ -2051,12 +2816,12 @@ mod tests { .await .expect("resume"); wait_until( - || fresh_log.lock().expect("lock").len() == 2, + || outbound_data_frames(&fresh_log).len() == 2, "buffered post-failure frames flush after resume", ) .await; assert_eq!( - *fresh_log.lock().expect("lock"), + outbound_data_frames(&fresh_log), vec!["kept1".to_owned(), "kept2".to_owned()], "post-failure frames survive; only the in-flight 'lost' frame is gone" ); @@ -2070,13 +2835,16 @@ mod tests { let (ctl_tx, ctl_rx) = mpsc::channel::(2); let (stop_tx, stop_rx) = mpsc::channel::<()>(1); let write_error = idle_write_error_slot(); + let (_prio_tx, prio_rx) = mpsc::channel::(4); let writer = tokio::spawn(run_writer( sink, out_rx, + prio_rx, ctl_rx, stop_rx, - TEST_PING_NEVER, + None, write_error.clone(), + None, )); ctl_tx.send(WriterControl::Pause).await.expect("pause"); *write_error.lock() = Some("frame send failed: stale broken pipe".to_owned()); @@ -2098,13 +2866,16 @@ mod tests { let (_out_tx, out_rx) = mpsc::channel::(4); let (_ctl_tx, ctl_rx) = mpsc::channel::(2); let (stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); let writer = tokio::spawn(run_writer( sink, out_rx, + prio_rx, ctl_rx, stop_rx, - TEST_PING_NEVER, + None, idle_write_error_slot(), + None, )); stop_tx.send(()).await.expect("stop"); tokio::time::timeout(Duration::from_secs(2), writer) @@ -2118,13 +2889,16 @@ mod tests { let (out_tx, out_rx) = mpsc::channel::(4); let (_ctl_tx, ctl_rx) = mpsc::channel::(2); let (_stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); let writer = tokio::spawn(run_writer( sink, out_rx, + prio_rx, ctl_rx, stop_rx, - TEST_PING_NEVER, + None, idle_write_error_slot(), + None, )); drop(out_tx); tokio::time::timeout(Duration::from_secs(2), writer) @@ -2138,13 +2912,16 @@ mod tests { let (_out_tx, out_rx) = mpsc::channel::(4); let (ctl_tx, ctl_rx) = mpsc::channel::(2); let (_stop_tx, stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); let writer = tokio::spawn(run_writer( sink, out_rx, + prio_rx, ctl_rx, stop_rx, - TEST_PING_NEVER, + None, idle_write_error_slot(), + None, )); drop(ctl_tx); tokio::time::timeout(Duration::from_secs(2), writer) @@ -2378,6 +3155,7 @@ mod tests { &mut stop_rx, &mut reconnect_rx, Duration::from_secs(75), + &conn.inner.outbound_tx, ), ) .await @@ -2404,6 +3182,7 @@ mod tests { &mut stop_rx, &mut reconnect_rx, Duration::from_secs(75), + &conn.inner.outbound_tx, ), ) .await @@ -2598,21 +3377,21 @@ mod tests { fn test_inbound() -> (InboundTx, InboundRx) { futures::channel::mpsc::unbounded() } - /// A zero or unset liveness deadline resolves to 2.5× the effective - /// ping cadence; a positive override is honored verbatim. Mirrors the + /// A zero or unset liveness deadline resolves to `min(4× ping, 120s)`; + /// a positive override is honored verbatim. Mirrors the /// `resolve_ws_ping_interval` clamp semantics. #[test] fn resolve_ws_liveness_deadline_clamps_zero_and_unset_to_default() { let ping = Duration::from_secs(30); assert_eq!( resolve_ws_liveness_deadline(None, ping), - Duration::from_secs(75) + Duration::from_secs(120) ); assert_eq!( resolve_ws_liveness_deadline(Some(Duration::ZERO), ping), - Duration::from_secs(75) + Duration::from_secs(120) ); - let custom = Duration::from_secs(120); + let custom = Duration::from_secs(45); assert_eq!(resolve_ws_liveness_deadline(Some(custom), ping), custom); } /// The per-attempt reconnect budget tracks the liveness deadline above @@ -2635,7 +3414,12 @@ mod tests { fn resolve_ws_liveness_deadline_scales_with_ping_override() { assert_eq!( resolve_ws_liveness_deadline(None, Duration::from_secs(10)), - Duration::from_secs(25) + Duration::from_secs(40) + ); + assert_eq!( + resolve_ws_liveness_deadline(None, Duration::from_secs(60)), + Duration::from_secs(120), + "default liveness is capped below hub idle_timeout", ); } #[tokio::test(start_paused = true)] @@ -2652,6 +3436,7 @@ mod tests { &mut stop_rx, &mut reconnect_rx, liveness, + &conn.inner.outbound_tx, ) .await; assert!(matches!( @@ -2666,7 +3451,7 @@ mod tests { drop(inbound_tx); } #[tokio::test(start_paused = true)] - async fn reader_deadline_rearms_on_any_inbound_frame() { + async fn reader_deadline_rearms_on_rtt_proof_frames() { let (conn, _demux, _outbound_rx) = test_connection(); let (inbound_tx, mut inbound_rx) = test_inbound(); let (_stop_tx, mut stop_rx) = mpsc::channel::<()>(1); @@ -2678,20 +3463,21 @@ mod tests { &mut stop_rx, &mut reconnect_rx, liveness, + &conn.inner.outbound_tx, ); tokio::pin!(phase); let frames = [ Message::Pong(Vec::new().into()), - Message::Ping(Vec::new().into()), - Message::Text(r#"{"jsonrpc":"2.0","method":"noop","params":{}}"#.into()), + Message::Text(r#"{"method":"pong","ts_ms":1}"#.into()), Message::Pong(Vec::new().into()), + Message::Text(r#"{"method":"pong","ts_ms":2}"#.into()), ]; for frame in frames { tokio::time::advance(liveness * 3 / 4).await; inbound_tx.unbounded_send(Ok(frame)).expect("send frame"); assert!( futures::poll!(phase.as_mut()).is_pending(), - "phase must stay live while frames keep arriving" + "phase must stay live while RTT-proof frames keep arriving" ); } tokio::time::advance(liveness - Duration::from_millis(1)).await; @@ -2712,6 +3498,94 @@ mod tests { } } } + #[test] + fn classify_inbound_hub_ping_is_app_ping_not_data() { + let (conn, _demux, _outbound_rx) = test_connection(); + assert!( + matches!( + classify_inbound_text(&conn.inner, r#"{"method":"ping","ts_ms":1}"#), + InboundText::AppPing { .. } + ), + "hub app ping must classify as AppPing" + ); + assert!( + matches!( + classify_inbound_text(&conn.inner, r#"{"method":"pong","ts_ms":1}"#), + InboundText::AppPong + ), + "hub app pong must classify as AppPong" + ); + } + #[tokio::test(start_paused = true)] + async fn reader_deadline_ignores_inbound_only_hub_pings() { + let (conn, _demux, _outbound_rx) = test_connection(); + let (prio_tx, prio_rx) = mpsc::channel::(4); + drop(prio_rx); + let (inbound_tx, mut inbound_rx) = test_inbound(); + let (_stop_tx, mut stop_rx) = mpsc::channel::<()>(1); + let (_reconnect_tx, mut reconnect_rx) = mpsc::channel::<()>(1); + let liveness = Duration::from_secs(75); + let phase = run_reader_phase( + &conn.inner, + &mut inbound_rx, + &mut stop_rx, + &mut reconnect_rx, + liveness, + &prio_tx, + ); + tokio::pin!(phase); + assert!( + futures::poll!(phase.as_mut()).is_pending(), + "phase must start pending" + ); + let pong_dropped_before = crate::metrics::heartbeat_pong_dropped_count(); + for _ in 0..3 { + inbound_tx + .unbounded_send(Ok(Message::Text(r#"{"method":"ping","ts_ms":1}"#.into()))) + .expect("send hub ping"); + inbound_tx + .unbounded_send(Ok(Message::Ping(Vec::new().into()))) + .expect("send ws ping"); + assert!( + futures::poll!(phase.as_mut()).is_pending(), + "inbound-only pings must not kill early" + ); + } + tokio::time::advance(liveness * 3 / 4).await; + inbound_tx + .unbounded_send(Ok(Message::Text(r#"{"method":"ping","ts_ms":2}"#.into()))) + .expect("late hub ping"); + assert!( + futures::poll!(phase.as_mut()).is_pending(), + "late hub ping must not re-arm" + ); + tokio::time::advance(liveness / 4 - Duration::from_millis(1)).await; + assert!( + futures::poll!(phase.as_mut()).is_pending(), + "still inside the original liveness window" + ); + tokio::time::advance(Duration::from_millis(1)).await; + let mut exit = None; + for _ in 0..16 { + match futures::poll!(phase.as_mut()) { + std::task::Poll::Ready(e) => { + exit = Some(e); + break; + } + std::task::Poll::Pending => tokio::task::yield_now().await, + } + } + match exit.expect("deadline should fire on original L") { + ConnectedExit::SocketClosed(DisconnectCause::LivenessDeadline) => {} + ConnectedExit::SocketClosed(cause) => panic!("wrong cause {}", cause.label()), + ConnectedExit::Stop => panic!("stop"), + ConnectedExit::TerminalClose(code) => panic!("terminal {code}"), + } + assert!( + crate::metrics::heartbeat_pong_dropped_count() > pong_dropped_before, + "dropped priority rx must count heartbeat_pong_dropped" + ); + } #[tokio::test(start_paused = true)] async fn reader_deadline_huge_override_saturates_instead_of_panicking() { let (conn, _demux, _outbound_rx) = test_connection(); @@ -2724,6 +3598,7 @@ mod tests { &mut stop_rx, &mut reconnect_rx, Duration::MAX, + &conn.inner.outbound_tx, ); tokio::pin!(phase); inbound_tx @@ -2749,8 +3624,19 @@ mod tests { Poll::Ready(Ok(())) } fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { - if let Message::Ping(payload) = item { - let _ = self.inbound.unbounded_send(Ok(Message::Pong(payload))); + match item { + Message::Ping(payload) => { + let _ = self.inbound.unbounded_send(Ok(Message::Pong(payload))); + } + Message::Text(text) => { + if let Ok(v) = serde_json::from_str::(text.as_ref()) + && v.get("method").and_then(serde_json::Value::as_str) == Some("ping") + && let Ok(pong) = serde_json::to_string(&PongFrame::new(now_unix_millis())) + { + let _ = self.inbound.unbounded_send(Ok(Message::Text(pong.into()))); + } + } + _ => {} } Ok(()) } @@ -2776,15 +3662,18 @@ mod tests { let (_out_tx, out_rx) = mpsc::channel::(4); let (ctl_tx, ctl_rx) = mpsc::channel::>(2); let (writer_stop_tx, writer_stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); let writer = tokio::spawn(run_writer( PongEchoSink { inbound: inbound_tx.clone(), }, out_rx, + prio_rx, ctl_rx, writer_stop_rx, - ping, + Some(ping), idle_write_error_slot(), + None, )); let (_stop_tx, mut stop_rx) = mpsc::channel::<()>(1); let (_reconnect_tx, mut reconnect_rx) = mpsc::channel::<()>(1); @@ -2795,6 +3684,7 @@ mod tests { &mut stop_rx, &mut reconnect_rx, deadline, + &conn.inner.outbound_tx, ); tokio::pin!(phase); tokio::select! { @@ -2815,6 +3705,7 @@ mod tests { &mut stop_rx, &mut reconnect_rx, deadline, + &conn.inner.outbound_tx, ); tokio::pin!(phase); tokio::select! { @@ -2827,6 +3718,80 @@ mod tests { writer_stop_tx.send(()).await.expect("stop"); writer.await.expect("writer task joins"); } + struct AppPingOnlyEchoSink { + inbound: InboundTx, + } + impl futures::Sink for AppPingOnlyEchoSink { + type Error = std::io::Error; + fn poll_ready( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { + if let Message::Text(text) = item + && text.contains("\"method\":\"ping\"") + && text.contains("ts_ms") + { + let pong = serde_json::to_string(&PongFrame::new(1)).expect("pong"); + let _ = self.inbound.unbounded_send(Ok(Message::Text(pong.into()))); + } + Ok(()) + } + fn poll_flush( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + fn poll_close( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + } + #[tokio::test(start_paused = true)] + async fn app_ping_only_keeps_idle_connection_alive_without_ws_pongs() { + let ping = resolve_ws_ping_interval(None); + let deadline = resolve_ws_liveness_deadline(None, ping); + let (conn, _demux, _outbound_rx) = test_connection(); + let (inbound_tx, mut inbound_rx) = test_inbound(); + let (_out_tx, out_rx) = mpsc::channel::(4); + let (_ctl_tx, ctl_rx) = mpsc::channel::>(2); + let (writer_stop_tx, writer_stop_rx) = mpsc::channel::<()>(1); + let (_prio_tx, prio_rx) = mpsc::channel::(4); + let writer = tokio::spawn(run_writer( + AppPingOnlyEchoSink { + inbound: inbound_tx, + }, + out_rx, + prio_rx, + ctl_rx, + writer_stop_rx, + Some(ping), + idle_write_error_slot(), + None, + )); + let (_stop_tx, mut stop_rx) = mpsc::channel::<()>(1); + let (_reconnect_tx, mut reconnect_rx) = mpsc::channel::<()>(1); + let phase = run_reader_phase( + &conn.inner, + &mut inbound_rx, + &mut stop_rx, + &mut reconnect_rx, + deadline, + &conn.inner.outbound_tx, + ); + tokio::pin!(phase); + tokio::select! { + _ = phase.as_mut() => panic!("app-pong-only keepalive tripped liveness"), + _ = tokio::time::sleep(deadline * 4) => {} + } + writer_stop_tx.send(()).await.expect("stop"); + writer.await.expect("join"); + } #[tokio::test] async fn reader_phase_close_frame_classification_unchanged() { use tokio_tungstenite::tungstenite::protocol::CloseFrame; @@ -2847,6 +3812,7 @@ mod tests { &mut stop_rx, &mut reconnect_rx, Duration::from_secs(75), + &conn.inner.outbound_tx, ) .await; assert!(matches!(exit, ConnectedExit::TerminalClose(4100))); diff --git a/crates/common/xai-computer-hub-sdk/src/harness.rs b/crates/common/xai-computer-hub-sdk/src/harness.rs index c0a4e00..a5791e8 100644 --- a/crates/common/xai-computer-hub-sdk/src/harness.rs +++ b/crates/common/xai-computer-hub-sdk/src/harness.rs @@ -3076,7 +3076,9 @@ mod tests { "supported_protocol_versions": ["1.0.0"], }); let _ = socket.send(Message::Text(ack.to_string().into())).await; - while let Some(Ok(Message::Text(text))) = socket.recv().await { + // Ignore WS Ping/Pong/Close: keepalive fires immediately after hello. + while let Some(Ok(msg)) = socket.recv().await { + let Message::Text(text) = msg else { continue }; let Ok(value) = serde_json::from_str::(text.as_ref()) else { continue; }; diff --git a/crates/common/xai-computer-hub-sdk/src/metrics.rs b/crates/common/xai-computer-hub-sdk/src/metrics.rs index f837f70..0956a8d 100644 --- a/crates/common/xai-computer-hub-sdk/src/metrics.rs +++ b/crates/common/xai-computer-hub-sdk/src/metrics.rs @@ -433,6 +433,11 @@ mod inner { HEARTBEAT_PONG_DROPPED_TOTAL.inc(); } + #[cfg(test)] + pub(crate) fn heartbeat_pong_dropped_count() -> u64 { + HEARTBEAT_PONG_DROPPED_TOTAL.get() + } + pub(crate) fn cancel_applied() { CANCEL_APPLIED_TOTAL.inc(); } @@ -561,6 +566,10 @@ mod inner { #[cfg(not(feature = "metrics"))] mod inner { + #[cfg(test)] + static TEST_HEARTBEAT_PONG_DROPPED: std::sync::atomic::AtomicU64 = + std::sync::atomic::AtomicU64::new(0); + pub(crate) fn pool_connections_inc() {} pub(crate) fn pool_connections_dec() {} pub(crate) fn pool_evictions_inc() {} @@ -584,7 +593,15 @@ mod inner { pub(crate) fn writer_sink_send_error() {} pub(crate) fn reconnect_writer_resume() {} pub(crate) fn liveness_deadline_expired() {} - pub(crate) fn heartbeat_pong_dropped() {} + pub(crate) fn heartbeat_pong_dropped() { + #[cfg(test)] + TEST_HEARTBEAT_PONG_DROPPED.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + + #[cfg(test)] + pub(crate) fn heartbeat_pong_dropped_count() -> u64 { + TEST_HEARTBEAT_PONG_DROPPED.load(std::sync::atomic::Ordering::Relaxed) + } pub(crate) fn cancel_applied() {} pub(crate) fn cancel_pending_tombstoned() {} pub(crate) fn cancel_no_target() {} @@ -619,6 +636,8 @@ pub(crate) use inner::demux_inbox_depth_set; pub(crate) use inner::disconnect_detail_class; pub(crate) use inner::early_notif_buffered; pub(crate) use inner::heartbeat_pong_dropped; +#[cfg(test)] +pub(crate) use inner::heartbeat_pong_dropped_count; pub(crate) use inner::hook_send; pub(crate) use inner::inbox_full_notification_dropped; pub(crate) use inner::inbox_full_reject_send_failed; diff --git a/crates/common/xai-computer-hub-sdk/src/server.rs b/crates/common/xai-computer-hub-sdk/src/server.rs index 0602b43..71ee0e2 100644 --- a/crates/common/xai-computer-hub-sdk/src/server.rs +++ b/crates/common/xai-computer-hub-sdk/src/server.rs @@ -294,18 +294,15 @@ impl ToolServerBuilder { } /// Override the inbound-liveness deadline on a freshly-opened - /// connection: if no inbound WebSocket frame of any kind arrives within - /// this window, the connection is declared dead and reconnected. This - /// catches silently dead transports (e.g. a VM snapshot restore or - /// NAT/LB flow expiry) that a send-only keepalive never notices. + /// connection: if no RTT proof (WS/app pong) arrives within this + /// window, the connection is declared dead and reconnected. + /// Hub→client pings and one-way data do not re-arm. /// - /// Default (also used for a zero value): 2.5× the effective ping - /// interval — 75s at the default 30s ping — which guarantees at least - /// two keepalive pings fit in every window, so a healthy-but-idle - /// connection (one pong per ping) can never trip it. Explicit values - /// are honored verbatim; keep them comfortably above the ping interval - /// for the same reason (a value at or below the ping interval churns - /// healthy idle connections and is logged as a warning at connect). + /// Default (also used for a zero value): `min(4× ping, 120s)` — 120s + /// at the default 30s ping, still under the hub's ~150s idle. Explicit + /// values are honored verbatim; keep them comfortably above the ping + /// interval (a value at or below the ping interval churns healthy idle + /// connections and is logged as a warning at connect). pub fn with_ws_liveness_deadline(mut self, deadline: std::time::Duration) -> Self { self.ws_liveness_deadline = Some(deadline); self @@ -398,8 +395,9 @@ impl ToolServerBuilder { self } - /// Optional callback fired once on the initial successful connect, before - /// the actor starts (so it happens-before any disconnect/reconnect). + /// Optional callback fired once on the initial successful connect, after + /// the writer task enters its loop and before the reader actor starts. + /// The first keepalive may still be in flight. pub fn on_connect(mut self, cb: F) -> Self where F: Fn() + Send + Sync + 'static, diff --git a/crates/common/xai-grok-compaction/src/code_compaction/failure.rs b/crates/common/xai-grok-compaction/src/code_compaction/failure.rs index b52bf60..031b714 100644 --- a/crates/common/xai-grok-compaction/src/code_compaction/failure.rs +++ b/crates/common/xai-grok-compaction/src/code_compaction/failure.rs @@ -33,6 +33,7 @@ pub fn is_context_length_error(message: &str) -> bool { || m.contains("maximum prompt length") || m.contains("maximum context length") || m.contains("context_length_exceeded") + || (m.contains("current message") && m.contains("exceeds budget")) } /// Classify an HTTP API failure (status + message) for the compaction retry @@ -183,6 +184,9 @@ mod tests { "exceeds the maximum prompt length", "This model's maximum context length is 128000 tokens", "error code: context_length_exceeded", + "Failed to start sampling: [conversation] Current message (1000000 tokens) exceeds budget (500000 tokens)", + "compact failed: API error (status 400 Bad Request): invalid-argument: Failed to start sampling: [conversation] Current message (1000000 tokens) exceeds budget (500000 tokens)", + "Current message (600000) exceeds budget (500000)", ] { assert!(is_context_length_error(msg), "should match: {msg}"); } @@ -190,6 +194,8 @@ mod tests { "internal server error", "rate limited", "connection reset by peer", + "Attached file content (300000 tokens) causes message to exceed budget", + "compact index estimate 2.0 GB exceeds budget 1.0 GB", ] { assert!(!is_context_length_error(msg), "should not match: {msg}"); } diff --git a/crates/common/xai-grok-compaction/src/code_compaction/sample.rs b/crates/common/xai-grok-compaction/src/code_compaction/sample.rs index 3f35188..3f5ace8 100644 --- a/crates/common/xai-grok-compaction/src/code_compaction/sample.rs +++ b/crates/common/xai-grok-compaction/src/code_compaction/sample.rs @@ -319,6 +319,26 @@ mod tests { assert_eq!(sampler.call_count(), 1, "overflow must not retry"); } + #[tokio::test] + async fn conversation_exceeds_budget_is_context_overflow() { + let sampler = + MockSampler::scripted(vec![Err(CompactionSampleError::Other(anyhow::anyhow!( + "API error (status 400 Bad Request): invalid-argument: \ + Failed to start sampling: [conversation] Current message \ + (1000000 tokens) exceeds budget (500000 tokens)" + )))]); + let err = run(&sampler, 3).await.expect_err("should fail"); + assert!(matches!( + err, + SampleRetryError::Failure { + deterministic: true, + context_overflow: true, + .. + } + )); + assert_eq!(sampler.call_count(), 1, "overflow must not retry"); + } + #[tokio::test] async fn transient_exhausted_is_non_deterministic_failure() { let sampler = MockSampler::scripted(vec![ diff --git a/crates/common/xai-grok-compaction/src/reminder.rs b/crates/common/xai-grok-compaction/src/reminder.rs index 68f1f02..c0e1000 100644 --- a/crates/common/xai-grok-compaction/src/reminder.rs +++ b/crates/common/xai-grok-compaction/src/reminder.rs @@ -16,6 +16,12 @@ //! sections (files, AGENTS.md, skills, MCP, memory). Callers pass **borrowed //! views** (`&str` over live state) so long fields (commands, todo content, //! descriptions, ids) are not cloned just to format. +//! +//! KEEP IN SYNC: the exact wording of these sections is a compatibility +//! surface — downstream mirrors reproduce it verbatim (grep for +//! `format_section_running_subagents` / `format_section_background_tasks` +//! and `section_todo_list` mirrors). Update them when changing any wording +//! here. // --------------------------------------------------------------------------- // Borrowed views over harness live state (no long-string clones) diff --git a/crates/common/xai-tool-protocol/src/frames.rs b/crates/common/xai-tool-protocol/src/frames.rs index 3b63e54..883c950 100644 --- a/crates/common/xai-tool-protocol/src/frames.rs +++ b/crates/common/xai-tool-protocol/src/frames.rs @@ -905,6 +905,66 @@ pub struct ToolServerStatusPayload { /// tasks. #[serde(default)] pub idle_ignores_background: bool, + /// Why `idle_since_ms` is being withheld, when it is. `None` ⇒ not withheld + /// (or the tool server is genuinely busy, which `active_tool_calls` + /// reports). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub withhold_reason: Option, + /// Epoch ms the current hold is measured from. Never `None` while + /// `withhold_reason` is `Some`. + /// + /// NOT the instant the withhold began, for the preview reasons. It is the + /// **real-use anchor**: the last routed request or tool call, floored at + /// process start. A status poll never moves it, which is the point — a + /// ceiling has to be measured from genuine use, or the pane could hold a + /// sandbox open forever by resetting the clock it is judged against. So + /// for a poll-pinned session this is *older* than the poll-only period, + /// and `now - withhold_since_ms` is time-since-real-use, not + /// time-spent-withholding. `Durability` is the exception: there it is the + /// true busy-since stamp. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub withhold_since_ms: Option, + /// `true` once the current hold has crossed its effective ceiling. The + /// verdict is published rather than re-derived because the ceiling is + /// per-session config only the sender can see. Always `false` today: no + /// ceilings are configured yet. + #[serde(default)] + pub withhold_capped: bool, + /// Open WebSocket (HMR) tunnels through the in-sandbox preview proxy. + /// Nonzero ⇒ a client is attached even if the preview is otherwise silent. + #[serde(default)] + pub preview_ws_tunnels_open: u32, +} + +/// Why a tool server is withholding `idle_since_ms`, ordered by strength of +/// evidence that someone is really using the sandbox. +/// +/// Carries an [`Unknown`](Self::Unknown) escape so adding a variant cannot +/// break older readers: without it, one unrecognised string would fail +/// deserialization of the **entire** status frame, taking the idle verdict down +/// with it. Sandbox binaries are pinned per session, so old and new report side +/// by side for at least a full session TTL. +/// +/// Deliberately not `#[non_exhaustive]`: in-workspace matches should stay +/// exhaustive so a new variant is a compile error at every decision site, which +/// is a different problem from wire tolerance. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum IdleWithholdReason { + /// Artifact producers or queued uploads outstanding. Already bounded by the + /// durability idle-hold cap. + Durability, + /// An open WebSocket tunnel or a routed request in flight. Never a status + /// poll, however long it is held. + PreviewAttached, + /// Recent `Routed` preview traffic — a human loading the app. + PreviewRouted, + /// Only the preview pane's own `/__grok-preview/status` liveness poll. + PreviewStatusOnly, + /// A reason this build does not recognise — a newer sender. Never + /// constructed locally; only produced by deserialization. + #[serde(other)] + Unknown, } impl ToolServerStatusPayload { @@ -1336,6 +1396,10 @@ mod tests { drain_started_ms: Some(1721234599999), turn_active: true, idle_ignores_background: false, + withhold_reason: None, + withhold_since_ms: None, + withhold_capped: false, + preview_ws_tunnels_open: 0, }; let json = serde_json::to_value(&payload).expect("serialize"); assert_eq!(json["upload_queue_pending"], 7); @@ -1350,6 +1414,81 @@ mod tests { assert_eq!(back, payload); } + /// `withhold_reason` is snake_case on the wire so it can be used directly + /// as a metric label. + #[test] + fn tool_server_status_payload_carries_withhold_fields() { + let payload = super::ToolServerStatusPayload { + status: super::ToolServerLifecycleStatus::Ready, + session_id: Some(sid()), + idle_since_ms: None, + withhold_reason: Some(super::IdleWithholdReason::PreviewStatusOnly), + withhold_since_ms: Some(1721234560000), + withhold_capped: true, + preview_ws_tunnels_open: 2, + ..Default::default() + }; + let json = serde_json::to_value(&payload).expect("serialize"); + assert_eq!(json["withhold_reason"], "preview_status_only"); + assert_eq!(json["withhold_since_ms"], 1721234560000u64); + assert_eq!(json["withhold_capped"], true); + assert_eq!(json["preview_ws_tunnels_open"], 2); + let back: super::ToolServerStatusPayload = + serde_json::from_value(json).expect("deserialize"); + assert_eq!(back, payload); + } + + /// An unrecognised reason from a newer sender must degrade to `Unknown`, + /// not fail the whole frame and take the idle verdict with it. + #[test] + fn unknown_withhold_reason_does_not_fail_the_frame() { + let json = serde_json::json!({ + "status": "ready", + "active_tool_calls": 0, + "background_tasks": 0, + "pending_tool_calls": 0, + "last_tool_call_started_ms": 0, + "last_tool_call_completed_ms": 0, + "uptime_ms": 1000, + "withhold_reason": "some_future_reason", + "withhold_since_ms": 1721234560000u64, + }); + let back: super::ToolServerStatusPayload = + serde_json::from_value(json).expect("a newer reason must not break the frame"); + assert_eq!( + back.withhold_reason, + Some(super::IdleWithholdReason::Unknown) + ); + assert_eq!( + back.withhold_since_ms, + Some(1721234560000), + "the rest of the frame must survive intact" + ); + } + + /// The fleet is version-pinned per session, so old and new binaries report + /// side by side for at least a full session TTL. + #[test] + fn tool_server_status_payload_withhold_fields_are_optional_on_the_wire() { + let json = serde_json::json!({ + "status": "ready", + "active_tool_calls": 0, + "background_tasks": 0, + "pending_tool_calls": 0, + "last_tool_call_started_ms": 0, + "last_tool_call_completed_ms": 0, + "uptime_ms": 1000, + }); + let back: super::ToolServerStatusPayload = + serde_json::from_value(json).expect("old payload must still deserialize"); + assert_eq!(back.withhold_reason, None); + assert_eq!(back.withhold_since_ms, None); + assert!(!back.withhold_capped); + assert_eq!(back.preview_ws_tunnels_open, 0); + // An absent reason is indistinguishable from "nothing is withheld" — + // what the field-coverage ratio exists to measure. + } + /// A legacy payload without the new fields deserializes with defaults. #[test] fn tool_server_status_payload_legacy_without_pr9_fields_defaults() { diff --git a/crates/common/xai-tool-protocol/src/lib.rs b/crates/common/xai-tool-protocol/src/lib.rs index 3061b1a..33fd674 100644 --- a/crates/common/xai-tool-protocol/src/lib.rs +++ b/crates/common/xai-tool-protocol/src/lib.rs @@ -38,21 +38,21 @@ pub use error_codes::{ }; pub use error_wire::ToolErrorWire; pub use frames::{ - AttachRoute, HookFrame, HookReplyFrame, LastSeq, LogsDonateParams, MAX_DONATION_BYTES, - MAX_LOG_RECORDS_PER_DONATION, MAX_METRICS_PER_DONATION, MAX_SPANS_PER_DONATION, - MAX_SYSTEM_NOTIFY_PAYLOAD_BYTES, MetricsDonateParams, NotificationFilter, PingFrame, PongFrame, - ServeParams, ServeResult, ServerBindAck, ServerBindOutcome, ServerBindParams, ServerInfo, - ServerUnbindAck, ServerUnbindOutcome, ServerUnbindParams, ServersListParams, ServersListResult, - SessionAttachServerParams, SessionAttachServerResult, SessionBindParams, SessionBindResult, - SessionBindServerParams, SessionBindServerResult, SessionCloseParams, SessionOpenParams, - SessionOpenResult, SessionUnbindParams, SessionUnbindServerParams, SubscribeAck, - SubscribeNotificationsParams, SubscribeOutcome, SystemNotifyParams, ToolCallParams, - ToolCallProgressFrame, ToolCallResult, ToolNotificationFrame, ToolSearchResult, - ToolServerConnectionStatus, ToolServerDisconnectReason, ToolServerEvictParams, - ToolServerGetStatusParams, ToolServerGetStatusResult, ToolServerLifecycleStatus, - ToolServerStatusPayload, ToolsChanged, ToolsListParams, ToolsListResult, ToolsSearchParams, - ToolsSearchResultBody, TracesDonateParams, UnsubscribeAck, UnsubscribeNotificationsParams, - UnsubscribeOutcome, + AttachRoute, HookFrame, HookReplyFrame, IdleWithholdReason, LastSeq, LogsDonateParams, + MAX_DONATION_BYTES, MAX_LOG_RECORDS_PER_DONATION, MAX_METRICS_PER_DONATION, + MAX_SPANS_PER_DONATION, MAX_SYSTEM_NOTIFY_PAYLOAD_BYTES, MetricsDonateParams, + NotificationFilter, PingFrame, PongFrame, ServeParams, ServeResult, ServerBindAck, + ServerBindOutcome, ServerBindParams, ServerInfo, ServerUnbindAck, ServerUnbindOutcome, + ServerUnbindParams, ServersListParams, ServersListResult, SessionAttachServerParams, + SessionAttachServerResult, SessionBindParams, SessionBindResult, SessionBindServerParams, + SessionBindServerResult, SessionCloseParams, SessionOpenParams, SessionOpenResult, + SessionUnbindParams, SessionUnbindServerParams, SubscribeAck, SubscribeNotificationsParams, + SubscribeOutcome, SystemNotifyParams, ToolCallParams, ToolCallProgressFrame, ToolCallResult, + ToolNotificationFrame, ToolSearchResult, ToolServerConnectionStatus, + ToolServerDisconnectReason, ToolServerEvictParams, ToolServerGetStatusParams, + ToolServerGetStatusResult, ToolServerLifecycleStatus, ToolServerStatusPayload, ToolsChanged, + ToolsListParams, ToolsListResult, ToolsSearchParams, ToolsSearchResultBody, TracesDonateParams, + UnsubscribeAck, UnsubscribeNotificationsParams, UnsubscribeOutcome, }; pub use handshake::{HelloAckMsg, HelloMsg, PROTOCOL_VERSION}; pub use hook::HookEvent; diff --git a/crates/common/xai-tool-types/src/lib.rs b/crates/common/xai-tool-types/src/lib.rs index 4da01ee..fd19d74 100644 --- a/crates/common/xai-tool-types/src/lib.rs +++ b/crates/common/xai-tool-types/src/lib.rs @@ -13,15 +13,16 @@ pub use serde_lenient::{ pub use task::{ BUILTIN_SUBAGENTS, BuiltinSubagent, EXPLORE_PROMPT, EXPLORE_SUBAGENT, GENERAL_PURPOSE_PROMPT, GENERAL_PURPOSE_SUBAGENT, KillTaskOutput, KillTaskResult, KillTaskToolInput, - KillTaskToolNaming, MAX_MULTI_WAIT_IDS, MultiTaskOutputResult, PLAN_PROMPT, PLAN_SUBAGENT, - SubagentCapabilityMode, SubagentCompletedOutput, SubagentDescriptor, SubagentIsolationMode, - SubagentToolNaming, TaskOutputOutput, TaskOutputResult, TaskOutputToolInput, - TaskOutputToolNaming, TaskToolInput, TaskToolNaming, WaitMode, WaitTasksToolInput, - WaitTasksToolNaming, build_kill_task_description, build_task_description, - build_task_output_description, build_wait_tasks_description, builtin_subagent_by_name, - default_subagent_type, format_resume_footer, format_subagent_completed, - format_subagent_started_background, is_not_sentinel, resolve_task_ids, sanitize_optional_arg, - task_output_waits, task_output_waits_from_json, + KillTaskToolNaming, MAX_MULTI_WAIT_IDS, MAX_WAIT_BLOCK_MS_DEFAULT, MAX_WAIT_MS_PLACEHOLDER, + MultiTaskOutputResult, PLAN_PROMPT, PLAN_SUBAGENT, SubagentCapabilityMode, + SubagentCompletedOutput, SubagentDescriptor, SubagentIsolationMode, SubagentToolNaming, + TaskOutputOutput, TaskOutputResult, TaskOutputToolInput, TaskOutputToolNaming, TaskToolInput, + TaskToolNaming, WaitMode, WaitTasksToolInput, WaitTasksToolNaming, build_kill_task_description, + build_task_description, build_task_output_description, build_wait_tasks_description, + builtin_subagent_by_name, default_subagent_type, format_resume_footer, + format_subagent_completed, format_subagent_started_background, format_wait_cap_ms, + is_not_sentinel, max_wait_block_ms, resolve_task_ids, sanitize_optional_arg, task_output_waits, + task_output_waits_from_json, }; pub use types::{ ArgumentType, SchemaType, ToolArgument, ToolDescription, ValidationError, ValidationErrors, diff --git a/crates/common/xai-tool-types/src/task.rs b/crates/common/xai-tool-types/src/task.rs index ffa0f74..4454b7c 100644 --- a/crates/common/xai-tool-types/src/task.rs +++ b/crates/common/xai-tool-types/src/task.rs @@ -340,8 +340,12 @@ pub struct TaskOutputToolInput { pub task_ids: Vec, /// When set and positive, wait up to this many milliseconds; omit or `0` polls. + /// + /// `{max_wait_ms}` is resolved at finalize from the session's wait ceiling, + /// which also pins it as the schema `maximum` — the tool description cannot + /// carry the bound alone, since randomization may replace it wholesale. #[schemars( - description = "Max wait time in milliseconds. A positive value waits for completion; omit or pass 0 for a non-blocking status poll." + description = "Max wait time in milliseconds, up to {max_wait_ms}. A positive value waits for completion; omit or pass 0 for a non-blocking status poll." )] #[serde(default)] pub timeout_ms: Option, @@ -381,6 +385,44 @@ pub fn task_output_waits(timeout_ms: Option) -> bool { timeout_ms.is_some_and(|ms| ms > 0) } +/// Default ceiling on a single blocking wait (`get_task_output` with a positive +/// `timeout_ms`, `wait_tasks`). Capping is safe because a completed task pings +/// the model, so a truncated wait costs one more poll, not the result. +pub const MAX_WAIT_BLOCK_MS_DEFAULT: u64 = 600_000; + +/// The blocking-wait ceiling in effect, honoring `GROK_MAX_WAIT_BLOCK_MS`. +/// +/// A host whose transport deadline is shorter than the default sets the env var +/// so the server enforces — and the tool descriptions advertise — the same +/// number the caller will actually wait for. Without that, a model believing the +/// default asks for a wait its own client will abandon first. +pub fn max_wait_block_ms() -> u64 { + std::env::var("GROK_MAX_WAIT_BLOCK_MS") + .ok() + .and_then(|raw| raw.parse::().ok()) + .unwrap_or(MAX_WAIT_BLOCK_MS_DEFAULT) +} + +/// Render a wait ceiling for tool descriptions, e.g. `600000 (~10 min)`. +/// +/// The unit is derived from the value, so it cannot drift from the millisecond +/// figure beside it. Both branches round *down*: a cap must never read as +/// longer than it is. +pub fn format_wait_cap_ms(ms: u64) -> String { + if ms < 60_000 { + format!("{ms} (~{} s)", ms / 1_000) + } else { + format!("{ms} (~{} min)", ms / 60_000) + } +} + +/// Placeholder the description builders emit for the wait ceiling. +/// +/// Resolved per session by `TruncationConfig::interpolate_description` in the +/// finalize loop, the same way `{max_lines_read}` is: the cap is client +/// configurable, so it cannot be baked in when the description is built. +pub const MAX_WAIT_MS_PLACEHOLDER: &str = "{max_wait_ms}"; + /// Same as [`task_output_waits`], from raw tool-arg JSON (fingerprint / doom-loop). pub fn task_output_waits_from_json(args: &serde_json::Value) -> bool { let timeout_ms = args.get("timeout_ms").and_then(|v| { @@ -504,7 +546,9 @@ pub struct WaitTasksToolInput { )] pub mode: WaitMode, - #[schemars(description = "Max wait time in milliseconds")] + /// Carries the same `{max_wait_ms}` marker as `TaskOutputToolInput`: this + /// tool blocks on the same ceiling, so it needs the same resolved bound. + #[schemars(description = "Max wait time in milliseconds, up to {max_wait_ms}")] #[serde(default)] pub timeout_ms: Option, } @@ -1023,12 +1067,13 @@ pub fn build_task_output_description(naming: &TaskOutputToolNaming) -> String { Some(r) => format!("\n- If output is large, use {r} on the output_file path"), None => String::new(), }; + let wait_cap = MAX_WAIT_MS_PLACEHOLDER; format!( "Get output and status from a background task{target_suffix}.\n\n\ Usage notes:\n\ - Pass {task_ids_param} with one or more ids from {sources}{monitor_note}; for a single task use a one-element array. Multiple ids with a positive {timeout_ms_param} wait until all complete\n\ - - Omit {timeout_ms_param} or pass 0 for a non-blocking status snapshot; set a positive {timeout_ms_param} to wait up to that many milliseconds, capped at ~10 min\n\ + - Omit {timeout_ms_param} or pass 0 for a non-blocking status snapshot; set a positive {timeout_ms_param} to wait up to that many milliseconds, capped at {wait_cap}\n\ - Returns current output, status, and exit code if completed{read_note}" ) } @@ -1061,13 +1106,15 @@ pub fn build_wait_tasks_description(naming: &WaitTasksToolNaming) -> String { (None, None) => "background tasks".to_string(), }; + let wait_cap = MAX_WAIT_MS_PLACEHOLDER; + format!( "Wait for multiple background tasks or subagents to complete.\n\n\ Prefer {background_retrieval_tool} with task_ids and a positive timeout_ms. This tool is kept for compatibility.\n\n\ Usage notes:\n\ - task_ids: list of task IDs from {sources}\n\ - mode: 'wait_all' or 'wait_any'\n\ - - timeout_ms: optional max wait, default 30s, capped at ~10 min" + - timeout_ms: optional max wait, default 30s, capped at {wait_cap}" ) } @@ -1528,6 +1575,19 @@ mod tests { ); } + #[test] + fn format_wait_cap_ms_derives_its_unit_and_rounds_down() { + assert_eq!( + format_wait_cap_ms(MAX_WAIT_BLOCK_MS_DEFAULT), + "600000 (~10 min)" + ); + assert_eq!(format_wait_cap_ms(300_000), "300000 (~5 min)"); + // Rounds down: 1.5 min must not read as 2. + assert_eq!(format_wait_cap_ms(90_000), "90000 (~1 min)"); + // Sub-minute caps switch unit rather than rendering "~0 min". + assert_eq!(format_wait_cap_ms(30_000), "30000 (~30 s)"); + } + #[test] fn task_output_description_tracks_renamed_params() { let desc = build_task_output_description(&TaskOutputToolNaming { @@ -1573,7 +1633,7 @@ mod tests { "Get output and status from a background task, monitor, or subagent.\n\n\ Usage notes:\n\ - Pass task_ids with one or more ids from background=true commands or subagents (a monitor's task_id is returned by monitor); for a single task use a one-element array. Multiple ids with a positive timeout_ms wait until all complete\n\ - - Omit timeout_ms or pass 0 for a non-blocking status snapshot; set a positive timeout_ms to wait up to that many milliseconds, capped at ~10 min\n\ + - Omit timeout_ms or pass 0 for a non-blocking status snapshot; set a positive timeout_ms to wait up to that many milliseconds, capped at {max_wait_ms}\n\ - Returns current output, status, and exit code if completed\n\ - If output is large, use read_file on the output_file path" ); @@ -1595,7 +1655,7 @@ mod tests { "Get output and status from a background task or subagent.\n\n\ Usage notes:\n\ - Pass task_ids with one or more ids from run_in_background=true subagents; for a single task use a one-element array. Multiple ids with a positive timeout_ms wait until all complete\n\ - - Omit timeout_ms or pass 0 for a non-blocking status snapshot; set a positive timeout_ms to wait up to that many milliseconds, capped at ~10 min\n\ + - Omit timeout_ms or pass 0 for a non-blocking status snapshot; set a positive timeout_ms to wait up to that many milliseconds, capped at {max_wait_ms}\n\ - Returns current output, status, and exit code if completed\n\ - If output is large, use read_file on the output_file path" ); @@ -1615,7 +1675,7 @@ mod tests { Usage notes:\n\ - task_ids: list of task IDs from background=true commands or subagents\n\ - mode: 'wait_all' or 'wait_any'\n\ - - timeout_ms: optional max wait, default 30s, capped at ~10 min" + - timeout_ms: optional max wait, default 30s, capped at {max_wait_ms}" ); } diff --git a/prod/mc/cli-chat-proxy-types/src/team_managed_config_types.rs b/prod/mc/cli-chat-proxy-types/src/team_managed_config_types.rs index e5b3fb1..eb4468d 100644 --- a/prod/mc/cli-chat-proxy-types/src/team_managed_config_types.rs +++ b/prod/mc/cli-chat-proxy-types/src/team_managed_config_types.rs @@ -2,7 +2,19 @@ //! the documents written here are served to CLIs by `/v1/deployment/config`. use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Deserializer, Serialize}; + +/// Deserialize a present field into `Some(_)` even when its value is `null`. +/// Plain `Option>` cannot distinguish the two: serde maps an explicit +/// `null` to `None`, the same as an absent field. Paired with `#[serde(default)]` +/// this gives the three states a partial update needs. +fn present_or_null<'de, T, D>(deserializer: D) -> Result>, D::Error> +where + T: Deserialize<'de>, + D: Deserializer<'de>, +{ + Option::deserialize(deserializer).map(Some) +} #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] pub struct TeamManagedConfig { @@ -16,16 +28,31 @@ pub struct TeamManagedConfig { pub updated_at: Option>, } -/// Full-replace: an absent field clears its stored document, so to update an -/// existing config echo both documents and `updated_at` from a GET — otherwise -/// a partial write silently drops the other (including `fail_closed`). +/// A partial update. Each document field is a double option, which is what lets +/// JSON say all three things: +/// +/// - absent (`None`) — leave the stored document alone +/// - `null` (`Some(None)`) — clear it +/// - a string (`Some(Some(_))`) — replace it +/// +/// An omitted field can therefore never erase a document, which is how an +/// editor that does not round-trip `requirements` stops being able to wipe a +/// team's enforced policy. A body that names neither document is a 400. #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] -#[serde(deny_unknown_fields)] // a typoed guard key must 400, not silently unguard +#[serde(deny_unknown_fields)] // a typoed key must 400, not silently leave a field alone pub struct SetTeamManagedConfigRequest { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub managed_config: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub requirements: Option, + #[serde( + default, + deserialize_with = "present_or_null", + skip_serializing_if = "Option::is_none" + )] + pub managed_config: Option>, + #[serde( + default, + deserialize_with = "present_or_null", + skip_serializing_if = "Option::is_none" + )] + pub requirements: Option>, /// Guard: the write fails 412 unless the stored `updated_at` equals this. #[serde(default, skip_serializing_if = "Option::is_none")] pub expected_updated_at: Option>, @@ -71,4 +98,42 @@ mod tests { let del: DeleteTeamManagedConfigRequest = serde_json::from_str("{}").unwrap(); assert_eq!(del.expected_updated_at, None); } + + /// The three states a partial update needs. Plain `Option>` + /// silently collapses `null` onto absent, which would make "clear this + /// document" unexpressible — hence `present_or_null`. + #[test] + fn set_request_distinguishes_absent_null_and_value() { + let absent: SetTeamManagedConfigRequest = + serde_json::from_str(r#"{"requirements":"fail_closed = true\n"}"#).unwrap(); + assert_eq!(absent.managed_config, None, "absent must not mean clear"); + assert_eq!( + absent.requirements, + Some(Some("fail_closed = true\n".to_owned())) + ); + + let cleared: SetTeamManagedConfigRequest = + serde_json::from_str(r#"{"managed_config":null}"#).unwrap(); + assert_eq!(cleared.managed_config, Some(None), "null must mean clear"); + assert_eq!(cleared.requirements, None); + + // Round-trips: absent stays absent, `null` stays `null`. + assert_eq!( + serde_json::to_string(&cleared).unwrap(), + r#"{"managed_config":null}"# + ); + assert_eq!( + serde_json::from_str::( + &serde_json::to_string(&cleared).unwrap() + ) + .unwrap(), + cleared + ); + + // A typoed key is still a 400 rather than a silent no-op. + assert!( + serde_json::from_str::(r#"{"managed_confgi":"x"}"#) + .is_err() + ); + } }