From 7a33cc7e565eb2e17225c72aff23d65b9cf4a86c Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Wed, 2 Sep 2026 17:09:08 -0500 Subject: [PATCH] Improve agent provider resilience --- AGENTS.md | 2 +- GALAXY.md | 18 +- REVIEW.md | 38 +-- app/src/ai/acp/transport.rs | 1 + .../agent/api/convert_conversation_tests.rs | 2 +- app/src/ai/agent/conversation.rs | 22 +- app/src/ai/bedrock/mod.rs | 23 -- app/src/ai/bedrock/models.rs | 50 ---- app/src/ai/blocklist/action_model.rs | 76 +++++- app/src/ai/blocklist/action_model_tests.rs | 20 ++ app/src/ai/blocklist/block.rs | 6 +- app/src/ai/blocklist/controller.rs | 63 ++++- .../blocklist/controller/response_stream.rs | 12 +- app/src/ai/blocklist/controller_tests.rs | 1 + .../ai/blocklist/usage/context_window_view.rs | 2 +- app/src/ai/crosscheck/reviewer.rs | 8 +- app/src/ai/llms.rs | 22 +- app/src/ai/llms_tests.rs | 11 +- app/src/ai/mod.rs | 2 - app/src/ai/openai/client.rs | 1 - app/src/ai/openai/response_translator.rs | 4 +- app/src/ai/prompt_builder/mod.rs | 2 +- app/src/ai/prompt_builder/tests.rs | 2 +- app/src/ai/prompt_builder/tools.rs | 2 +- app/src/ai/{bedrock => provider}/client.rs | 1 - app/src/ai/{bedrock => provider}/convert.rs | 0 .../{bedrock => provider}/convert_request.rs | 0 .../ai/{bedrock => provider}/convert_tests.rs | 0 app/src/ai/{bedrock => provider}/crash_log.rs | 0 .../ai/{bedrock => provider}/diagnostic.rs | 0 app/src/ai/{bedrock => provider}/discovery.rs | 1 - app/src/ai/{bedrock => provider}/e2e_tests.rs | 3 +- .../{bedrock => provider}/external_config.rs | 1 - .../external_config_tests.rs | 0 .../integration_tests.rs | 1 - app/src/ai/provider/mod.rs | 46 +++- app/src/ai/provider/models.rs | 19 ++ .../ai/{bedrock => provider}/models_tests.rs | 31 --- .../request_translator.rs | 0 .../request_translator_tests.rs | 2 +- .../response_translator.rs | 0 .../response_translator_tests.rs | 0 .../ai/{bedrock => provider}/settings_view.rs | 4 +- .../test_fixtures/sample_project/Cargo.toml | 0 .../test_fixtures/sample_project/README.md | 0 .../test_fixtures/sample_project/src/lib.rs | 0 .../sample_project/src/lib_tests.rs | 0 .../test_fixtures/sample_project/src/main.rs | 0 .../test_fixtures/sample_project/src/utils.rs | 0 app/src/ai/runtime/event_translator.rs | 219 ++++++++++++++++-- app/src/ai/runtime/event_translator_tests.rs | 118 +++++++++- .../ai/runtime/provider_run_coordinator.rs | 76 ++++++ .../runtime/provider_run_coordinator_tests.rs | 44 ++++ app/src/ai/runtime/rig.rs | 8 +- app/src/ai/runtime/rig_request.rs | 34 ++- app/src/ai/runtime/rig_request_tests.rs | 1 - app/src/ai/runtime/rig_tests.rs | 1 - app/src/ai/runtime/rig_tool.rs | 2 +- app/src/code_review/code_review_header/mod.rs | 3 +- app/src/settings/ai.rs | 10 - app/src/settings/ai_tests.rs | 1 - app/src/settings_view/ai_page.rs | 6 - app/src/settings_view/provider_setup_view.rs | 5 +- app/src/terminal/input/slash_commands/mod.rs | 8 +- app/src/terminal/view.rs | 90 ++++--- crates/galaxy_agent_core/src/provider_run.rs | 23 ++ .../src/provider_run_tests.rs | 13 ++ crates/galaxy_agent_rig/src/stream.rs | 88 ++++++- crates/galaxy_agent_rig/src/stream_tests.rs | 55 +++++ crates/integration/src/test/rig_runtime.rs | 1 - migration-docs/project-architecture.md | 4 +- plans/galaxy-refactor.md | 8 +- 72 files changed, 997 insertions(+), 320 deletions(-) delete mode 100644 app/src/ai/bedrock/mod.rs delete mode 100644 app/src/ai/bedrock/models.rs rename app/src/ai/{bedrock => provider}/client.rs (99%) rename app/src/ai/{bedrock => provider}/convert.rs (100%) rename app/src/ai/{bedrock => provider}/convert_request.rs (100%) rename app/src/ai/{bedrock => provider}/convert_tests.rs (100%) rename app/src/ai/{bedrock => provider}/crash_log.rs (100%) rename app/src/ai/{bedrock => provider}/diagnostic.rs (100%) rename app/src/ai/{bedrock => provider}/discovery.rs (99%) rename app/src/ai/{bedrock => provider}/e2e_tests.rs (99%) rename app/src/ai/{bedrock => provider}/external_config.rs (99%) rename app/src/ai/{bedrock => provider}/external_config_tests.rs (100%) rename app/src/ai/{bedrock => provider}/integration_tests.rs (99%) create mode 100644 app/src/ai/provider/models.rs rename app/src/ai/{bedrock => provider}/models_tests.rs (74%) rename app/src/ai/{bedrock => provider}/request_translator.rs (100%) rename app/src/ai/{bedrock => provider}/request_translator_tests.rs (99%) rename app/src/ai/{bedrock => provider}/response_translator.rs (100%) rename app/src/ai/{bedrock => provider}/response_translator_tests.rs (100%) rename app/src/ai/{bedrock => provider}/settings_view.rs (98%) rename app/src/ai/{bedrock => provider}/test_fixtures/sample_project/Cargo.toml (100%) rename app/src/ai/{bedrock => provider}/test_fixtures/sample_project/README.md (100%) rename app/src/ai/{bedrock => provider}/test_fixtures/sample_project/src/lib.rs (100%) rename app/src/ai/{bedrock => provider}/test_fixtures/sample_project/src/lib_tests.rs (100%) rename app/src/ai/{bedrock => provider}/test_fixtures/sample_project/src/main.rs (100%) rename app/src/ai/{bedrock => provider}/test_fixtures/sample_project/src/utils.rs (100%) diff --git a/AGENTS.md b/AGENTS.md index 234d7210..088dd3e3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -78,7 +78,7 @@ ACP: controller.rs → ResponseStream → acp_output_stream (session-owned lifec - `types.rs` — `ConversationMessage`, `MessageRole`, `MessageContent`, `ContentPart`, `ToolDefinition` - `mod.rs` — `ProviderConfig` enum (Bedrock | OpenAI | None) -**Bedrock provider** in `app/src/ai/bedrock/`: +**Bedrock provider** in `app/src/ai/provider/`: - `runtime.rs` — Native one-call `AgentRuntime` over `ConverseStream`, including cancellation, reasoning signatures, token usage, and stop/error classification - `request_translator.rs` — Shared Bedrock message sanitization and tool definitions - `response_translator.rs` — Compatibility conversion helpers used by tests and background flows diff --git a/GALAXY.md b/GALAXY.md index 29553b4b..2e1d9c5e 100644 --- a/GALAXY.md +++ b/GALAXY.md @@ -56,11 +56,11 @@ This is a Rust-based terminal emulator with a custom UI framework called **Galax ### AI Architecture Galaxy uses **Amazon Bedrock** as the sole AI provider. The client calls Bedrock directly: -- `app/src/ai/bedrock/client.rs` - AWS Bedrock client (`BedrockClient::converse_stream`) -- `app/src/ai/bedrock/convert_request.rs` - System prompt construction, message/tool extraction -- `app/src/ai/bedrock/convert.rs` - Conversion to Bedrock wire format -- `app/src/ai/bedrock/tool_docs.rs` - Tool documentation (served via `get_tool_documentation` meta-tool) -- `app/src/ai/bedrock/stream.rs` - Response stream processing +- `app/src/ai/provider/client.rs` - AWS Bedrock client (`BedrockClient::converse_stream`) +- `app/src/ai/provider/convert_request.rs` - System prompt construction, message/tool extraction +- `app/src/ai/provider/convert.rs` - Conversion to Bedrock wire format +- `app/src/ai/provider/tool_docs.rs` - Tool documentation (served via `get_tool_documentation` meta-tool) +- `app/src/ai/provider/stream.rs` - Response stream processing System prompt is dynamically built from request context (OS, shell, pwd, git, project rules, global rules, skills, MCP servers). @@ -68,7 +68,7 @@ Global rules are loaded from `~/.galaxy-ai/rules/*.md` (filename = rule name, co Project rules are loaded from `GALAXY.md` or `AGENTS.md` files found in the project directory tree. -**Model Discovery** (`app/src/ai/bedrock/discovery.rs`): +**Model Discovery** (`app/src/ai/provider/discovery.rs`): At startup (and on manual refresh from the Bedrock settings page), Galaxy discovers available models using the user-selected auth settings (profile/SSO/static keys — no external config fallback) via: 1. **STS GetCallerIdentity** — validates AWS credentials before proceeding 2. **ListInferenceProfiles** (system-defined) + **ListFoundationModels** (TEXT output, ON_DEMAND) — fetched in parallel; results are deduplicated by underlying foundation model ID with inference profiles taking priority (they include cross-region routing) @@ -80,14 +80,14 @@ The settings page refresh (`RefreshAwsBedrock` in `app/src/settings_view/ai_page The `[1m]` suffix is an internal marker stripped by `strip_context_marker()` in `client.rs` before API calls. It's used by `context_window_for_model()` in `response_translator.rs` to report the correct context window size. -**External Config Fallback** (`app/src/ai/bedrock/external_config.rs`): +**External Config Fallback** (`app/src/ai/provider/external_config.rs`): When Galaxy's own Bedrock settings are at defaults, it falls back to configurations from: 1. **Claude Code** (`~/.claude/settings.json`) — reads `env.AWS_PROFILE`, `env.AWS_REGION`, and `env.DCP_MODEL_MAP` (ARN-based model mappings) 2. **OpenCode** (`~/.config/opencode/opencode.json`) — reads `provider.amazon-bedrock.options.profile` and `.region` Priority: Galaxy explicit settings > Claude Code > OpenCode > hardcoded defaults. Fallback only applies when profile is `"default"` (for profile) or empty (for region/models). External model ARNs are merged with Galaxy's built-in default model list. -**Token Usage & Cost Tracking** (`app/src/ai/bedrock/response_translator.rs`): +**Token Usage & Cost Tracking** (`app/src/ai/provider/response_translator.rs`): The Bedrock stream extracts full token metadata from responses: `input_tokens`, `output_tokens`, `cache_read_input_tokens`, `cache_write_input_tokens`. These flow through `build_stream_finished` → `TokenUsage` struct → `conversation.update_cost_and_usage_for_request()`. Cost is estimated per-model using Bedrock pricing. Displayed in: - **Agent management cards** — total token count in metadata row - **Conversation usage footer** — full breakdown (input/output/cache read/cache write) + estimated cost @@ -100,7 +100,7 @@ When context window usage reaches 85%, Galaxy automatically summarizes older mes - No UI shown — user only sees context usage drop - `recall_tool_history` tool lets the agent retrieve past tool outputs that were summarized away -**Failed Tool Call Visibility** (`app/src/ai/bedrock/response_translator.rs`): +**Failed Tool Call Visibility** (`app/src/ai/provider/response_translator.rs`): When the model calls an unknown/hallucinated tool name, the response translator now emits a visible `AgentOutput` text message to the UI (via `build_add_agent_output_message`) showing what tool was attempted and the error. Previously, synthetic error results were only stored in history (for Bedrock message ordering) but never rendered. **Loop Prevention Guardrail** (`app/src/ai/blocklist/controller.rs`): diff --git a/REVIEW.md b/REVIEW.md index aec464c6..d06b97e3 100644 --- a/REVIEW.md +++ b/REVIEW.md @@ -194,18 +194,18 @@ This is the primary new feature. A complete AWS Bedrock direct-call path runs al | File | Purpose | |---|---| -| `app/src/ai/bedrock/diagnostic.rs` | `BedrockDiagnosticLogger` — file-based diagnostic logging, activated via `GALAXY_BEDROCK_DIAGNOSTICS=1`, 10MB rotation | -| `app/src/ai/bedrock/e2e_tests.rs` | Comprehensive E2E integration tests (requires `BEDROCK_INTEGRATION_TEST` env var) | -| `app/src/ai/bedrock/integration_tests.rs` | Unit-level stream output collection tests | +| `app/src/ai/provider/diagnostic.rs` | `BedrockDiagnosticLogger` — file-based diagnostic logging, activated via `GALAXY_BEDROCK_DIAGNOSTICS=1`, 10MB rotation | +| `app/src/ai/provider/e2e_tests.rs` | Comprehensive E2E integration tests (requires `BEDROCK_INTEGRATION_TEST` env var) | +| `app/src/ai/provider/integration_tests.rs` | Unit-level stream output collection tests | ### Modified Files | File | Key Change | |---|---| -| `app/src/ai/bedrock/client.rs` | Added `needs_create_task: bool` and `diagnostic_logger` params to `converse_stream()` | -| `app/src/ai/bedrock/convert.rs` | Added `Clone`, `Debug`, `PartialEq` derives to core types | -| `app/src/ai/bedrock/convert_request.rs` | New input type handlers: `InitProjectRules`, `CreateEnvironment`, `CreateNewProject`, `CloneRepository`, `AutoCodeDiffQuery`, `ResumeConversation`, `QueryWithCannedResponse`, `CodeReview` | -| `app/src/ai/bedrock/stream.rs` | Buffered text flush threshold, `needs_create_task` / `CreateTask` event, reasoning content discarded | +| `app/src/ai/provider/client.rs` | Added `needs_create_task: bool` and `diagnostic_logger` params to `converse_stream()` | +| `app/src/ai/provider/convert.rs` | Added `Clone`, `Debug`, `PartialEq` derives to core types | +| `app/src/ai/provider/convert_request.rs` | New input type handlers: `InitProjectRules`, `CreateEnvironment`, `CreateNewProject`, `CloneRepository`, `AutoCodeDiffQuery`, `ResumeConversation`, `QueryWithCannedResponse`, `CodeReview` | +| `app/src/ai/provider/stream.rs` | Buffered text flush threshold, `needs_create_task` / `CreateTask` event, reasoning content discarded | | `app/src/ai/agent/api/impl.rs` | `needs_create_task` detection, diagnostic logger wiring, verbose debug logging | | `app/src/ai/llms.rs` | Bedrock models added to `coding` and `cli_agent` feature slots (previously only `agent_mode`) | @@ -220,26 +220,26 @@ This is the primary new feature. A complete AWS Bedrock direct-call path runs al **`needs_create_task` detection logic** (`app/src/ai/agent/api/impl.rs`) Detected via `task_context.tasks.is_empty()`. If the tasks list is populated with a placeholder or shell task entry before the first real task is submitted, this flag will be `false` and `CreateTask` will never be emitted, silently breaking the task creation flow for Bedrock on the first turn. -**E2E test tool_use_id mismatch** (`app/src/ai/bedrock/e2e_tests.rs`) +**E2E test tool_use_id mismatch** (`app/src/ai/provider/e2e_tests.rs`) In `test_agent_multi_turn_tool_use_produces_output`, the test synthesizes tool result IDs as `tool_{total_turns}`, but the actual tool call ID emitted from the stream is a different value. This means the tool_call_id round-trip assertion will fail or be skipped silently. The test may pass vacuously. -**`ensure_tool_results_paired()` synthesized results** (`app/src/ai/bedrock/convert_request.rs`) +**`ensure_tool_results_paired()` synthesized results** (`app/src/ai/provider/convert_request.rs`) This function inserts placeholder tool results for orphaned tool calls. The synthesized content is empty/default, which could confuse the model on subsequent turns. The synthesized result should include a meaningful message like `"Tool result not available"` rather than empty content. #### 🟡 Medium -**`TEXT_FLUSH_THRESHOLD = 20`** (`app/src/ai/bedrock/stream.rs`) +**`TEXT_FLUSH_THRESHOLD = 20`** (`app/src/ai/provider/stream.rs`) Text fragments shorter than 20 characters before a tool call boundary are silently discarded. This can drop model preamble text such as `"Let me look at that..."` or `"OK."`. Consider buffering until a natural boundary (tool call or stop event) rather than a character threshold, or surface buffered content even if short. -**`converse_stream()` wide parameter surface** (`app/src/ai/bedrock/client.rs`) +**`converse_stream()` wide parameter surface** (`app/src/ai/provider/client.rs`) The function now takes 9 parameters including `needs_create_task` in the middle of the list. This is fragile and hard to read at call sites. Recommend grouping into a `BedrockConversationConfig` struct. -**Diagnostic logger privacy concern** (`app/src/ai/bedrock/diagnostic.rs`) +**Diagnostic logger privacy concern** (`app/src/ai/provider/diagnostic.rs`) `log_protobuf_input` uses Rust's `{:?}` debug format rather than structured JSON. This is inconsistent with the rest of the logger (which emits JSON) and may dump sensitive content (full conversation history, file contents) to disk in an unstructured format. Consider either JSON serialization or explicit redaction. #### 🟢 Low -**Reasoning content silently discarded** (`app/src/ai/bedrock/stream.rs`) +**Reasoning content silently discarded** (`app/src/ai/provider/stream.rs`) Reasoning/thinking content from models that emit it (e.g., Claude 3.7 extended thinking) is currently discarded. This is an intentional product decision but should be documented in a comment so future contributors understand why. **`diagnostic.rs` log path** — Uses `warp_logging::log_directory()`. The log file will land in the `warp_logging`-configured directory. Verify this resolves to the correct Galaxy AI log directory on all platforms. @@ -316,7 +316,7 @@ All internal crates retain `warp_*` naming: |---|---|---| | C1 | `WARP_INTEGRATION` / `GALAXY_INTEGRATION` split — breaks integration test detection at runtime | `app/src/lib.rs:658` | | C2 | Windows TTY env vars completely skipped — Windows shell integration broken | `app/src/terminal/local_tty/windows/environment.rs` | -| C3 | E2E test tool_use_id mismatch — multi-turn tool use test may be vacuously passing | `app/src/ai/bedrock/e2e_tests.rs` | +| C3 | E2E test tool_use_id mismatch — multi-turn tool use test may be vacuously passing | `app/src/ai/provider/e2e_tests.rs` | ### 🟠 High — Fix Before Beta @@ -335,11 +335,11 @@ All internal crates retain `warp_*` naming: | # | Issue | File(s) | |---|---|---| | M1 | All `warp.dev` URLs untouched — wrong privacy policy, docs, and external links | `app/src/util/links.rs` + 7 other files | -| M2 | `converse_stream()` 9-param signature — should use config struct | `app/src/ai/bedrock/client.rs` | -| M3 | `TEXT_FLUSH_THRESHOLD = 20` may silently discard model preamble text | `app/src/ai/bedrock/stream.rs` | -| M4 | `ensure_tool_results_paired()` synthesizes empty tool results — may confuse model | `app/src/ai/bedrock/convert_request.rs` | +| M2 | `converse_stream()` 9-param signature — should use config struct | `app/src/ai/provider/client.rs` | +| M3 | `TEXT_FLUSH_THRESHOLD = 20` may silently discard model preamble text | `app/src/ai/provider/stream.rs` | +| M4 | `ensure_tool_results_paired()` synthesizes empty tool results — may confuse model | `app/src/ai/provider/convert_request.rs` | | M5 | `WARP_CHANNEL_VERSIONS_PATH` not renamed in autoupdate | `app/src/autoupdate/` | -| M6 | Diagnostic logger `log_protobuf_input` uses debug format, not JSON | `app/src/ai/bedrock/diagnostic.rs` | +| M6 | Diagnostic logger `log_protobuf_input` uses debug format, not JSON | `app/src/ai/provider/diagnostic.rs` | | M7 | `WARP_USER_SECRET` build-time env not renamed | `app/src/auth/auth_state.rs:107` | | M8 | `DEFAULT_UI_FONT_NAME = ""` sentinel is a code smell | `app/src/settings/font.rs` | | M9 | Referral theme display names still say `"Warp Referral"` | `app/src/themes/theme.rs` | @@ -351,7 +351,7 @@ All internal crates retain `warp_*` naming: | # | Issue | File(s) | |---|---|---| | L1 | Stale comment: *"Use the word 'Warp'..."* | `app/src/workspace/view.rs:546` | -| L2 | Reasoning content discard should be documented with a comment | `app/src/ai/bedrock/stream.rs` | +| L2 | Reasoning content discard should be documented with a comment | `app/src/ai/provider/stream.rs` | | L3 | `warp_*` crate naming — document intent (internal stable vs. rebrand needed) | All `Cargo.toml` files | | L4 | `apply_samsung_brand_preset()` full implementation not reviewed | `app/src/settings_view/appearance_page.rs` | | L5 | `WARP.md` build commands still reference `cargo bundle --bin warp` — should be updated | `WARP.md` | diff --git a/app/src/ai/acp/transport.rs b/app/src/ai/acp/transport.rs index a37a4b22..f9d24b36 100644 --- a/app/src/ai/acp/transport.rs +++ b/app/src/ai/acp/transport.rs @@ -228,6 +228,7 @@ fn response_translator( max_context_tokens: None, capabilities: RuntimeCapabilities::session_runtime(), empty_output_message: Some("> ACP agent completed without a text response.".to_owned()), + todo_items: None, }) } diff --git a/app/src/ai/agent/api/convert_conversation_tests.rs b/app/src/ai/agent/api/convert_conversation_tests.rs index 59d1e26c..23ae568a 100644 --- a/app/src/ai/agent/api/convert_conversation_tests.rs +++ b/app/src/ai/agent/api/convert_conversation_tests.rs @@ -2145,7 +2145,7 @@ fn test_internal_command_completion_assessment_restores_output_without_visible_i }; mark_internal_command_completion_assessment(&mut hidden_assessment); let provider_history = - crate::ai::bedrock::request_translator::convert_proto_message(&hidden_assessment) + crate::ai::provider::request_translator::convert_proto_message(&hidden_assessment) .expect("hidden assessment should remain in provider history"); assert_eq!( provider_history.role, diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index 438d43e2..06b708da 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -334,8 +334,8 @@ pub struct AIConversation { /// Whether the user has pinned this child agent in the orchestration /// pill bar. Persisted via `AgentConversationData.pinned`. pinned: bool, - bedrock_message_history: Vec, - tool_result_archive: Vec, + bedrock_message_history: Vec, + tool_result_archive: Vec, progressive_summary: Option, messages_summarized_up_to: usize, current_context_tokens: u32, @@ -448,10 +448,10 @@ impl AIConversation { tasks: Vec, conversation_data: Option, ) -> Result { - let bedrock_message_history: Vec = tasks + let bedrock_message_history: Vec = tasks .iter() .flat_map(|task| task.messages.iter()) - .filter_map(crate::ai::bedrock::request_translator::convert_proto_message) + .filter_map(crate::ai::provider::request_translator::convert_proto_message) .collect(); let (task_store, todo_lists, status) = if tasks.is_empty() { @@ -800,32 +800,32 @@ impl AIConversation { self.has_pending_progressive_summary = val; } - pub fn bedrock_message_history(&self) -> &[crate::ai::bedrock::convert::ConversationMessage] { + pub fn bedrock_message_history(&self) -> &[crate::ai::provider::convert::ConversationMessage] { &self.bedrock_message_history } pub fn bedrock_message_history_mut( &mut self, - ) -> &mut Vec { + ) -> &mut Vec { &mut self.bedrock_message_history } - pub fn tool_result_archive(&self) -> &[crate::ai::bedrock::convert::ConversationMessage] { + pub fn tool_result_archive(&self) -> &[crate::ai::provider::convert::ConversationMessage] { &self.tool_result_archive } pub fn archive_tool_results( &mut self, - messages: Vec, + messages: Vec, ) { - use crate::ai::bedrock::convert::{ContentPart, MessageContent}; + use crate::ai::provider::convert::{ContentPart, MessageContent}; // Cap the archive to prevent unbounded growth. The archive is only used by // `recall_tool_history` which already truncates individual results to 50K chars, // so retaining the most recent entries is sufficient for lookup. const MAX_TOOL_RESULT_ARCHIVE_ENTRIES: usize = 400; - let mut pending_tool_uses: Vec = + let mut pending_tool_uses: Vec = Vec::new(); for msg in messages { @@ -873,7 +873,7 @@ impl AIConversation { pub fn append_to_bedrock_history( &mut self, - messages: Vec, + messages: Vec, ) { self.bedrock_message_history.extend(messages); } diff --git a/app/src/ai/bedrock/mod.rs b/app/src/ai/bedrock/mod.rs deleted file mode 100644 index d3b034ed..00000000 --- a/app/src/ai/bedrock/mod.rs +++ /dev/null @@ -1,23 +0,0 @@ -pub mod client; -pub mod convert; -pub mod crash_log; -pub mod diagnostic; -pub mod discovery; -pub mod external_config; -pub mod models; -pub mod request_translator; -pub mod response_translator; -pub mod settings_view; - -#[cfg(test)] -mod convert_tests; -#[cfg(test)] -#[allow(dead_code)] -mod e2e_tests; -#[cfg(test)] -#[allow(dead_code)] -mod integration_tests; -#[cfg(test)] -mod models_tests; -#[cfg(test)] -mod response_translator_tests; diff --git a/app/src/ai/bedrock/models.rs b/app/src/ai/bedrock/models.rs deleted file mode 100644 index af73d25b..00000000 --- a/app/src/ai/bedrock/models.rs +++ /dev/null @@ -1,50 +0,0 @@ -#![allow(dead_code)] - -use crate::settings::ai::BedrockModelConfig; - -pub fn configured_model_uses_rig( - selected_model_id: &str, - configured_models: &[BedrockModelConfig], - region: &str, - cross_region_inference: bool, -) -> bool { - let selected_model_id = strip_context_marker(selected_model_id); - configured_models.iter().any(|model| { - if !model.use_rig { - return false; - } - let configured_model_id = strip_context_marker(&model.model_id); - if configured_model_id == selected_model_id { - return true; - } - galaxy_agent_rig::resolve_bedrock_model_id(&model.model_id, region, cross_region_inference) - .is_ok_and(|resolved| strip_context_marker(&resolved) == selected_model_id) - }) -} - -fn strip_context_marker(model_id: &str) -> &str { - model_id - .strip_suffix("[1m]") - .or_else(|| model_id.strip_suffix("[1M]")) - .unwrap_or(model_id) -} - -pub fn apply_cross_region_prefix(model_id: &str, region: &str) -> String { - if model_id.starts_with("arn:") { - return model_id.to_string(); - } - - if model_id.contains('.') && model_id.split('.').next().unwrap_or("").len() <= 6 { - return model_id.to_string(); - } - - let prefix = match region { - r if r.starts_with("us-") || r.starts_with("ca-") => "us", - r if r.starts_with("eu-") || r == "il-central-1" => "eu", - r if r == "ap-northeast-1" || r == "ap-northeast-3" => "jp", - r if r == "ap-southeast-2" || r == "ap-southeast-4" || r == "ap-southeast-6" => "au", - r if r.starts_with("ap-") => "apac", - _ => return model_id.to_string(), - }; - format!("{}.{}", prefix, model_id) -} diff --git a/app/src/ai/blocklist/action_model.rs b/app/src/ai/blocklist/action_model.rs index 5bd88405..753a6c2d 100644 --- a/app/src/ai/blocklist/action_model.rs +++ b/app/src/ai/blocklist/action_model.rs @@ -1754,23 +1754,77 @@ impl BlocklistAIActionModel { batch: &PendingToolBatch, ctx: &mut ModelContext, ) -> Result<(), ProviderActionQueueError> { - let refs = provider_action_correlations(&actions, conversation_id, batch)?; - for ((_, action_id), _) in &refs { - if self - .provider_tool_executions - .contains_key(&(conversation_id, action_id.clone())) + let action_ids = self.provider_action_ids_to_enqueue(&actions, conversation_id, batch)?; + let refs = provider_action_correlations(&actions, conversation_id, batch)? + .into_iter() + .filter(|((_, action_id), _)| action_ids.contains(action_id)); + self.provider_tool_executions.extend(refs); + self.executor.update(ctx, |executor, ctx| { + executor.mark_restored_actions(conversation_id, &recovery_action_ids, ctx); + }); + let actions: Vec = actions + .into_iter() + .filter(|action| action_ids.contains(&action.id)) + .collect(); + if !actions.is_empty() { + self.queue_actions(actions, conversation_id, ctx); + } + Ok(()) + } + + pub(super) fn provider_action_ids_to_enqueue( + &self, + actions: &[AIAgentAction], + conversation_id: AIConversationId, + batch: &PendingToolBatch, + ) -> Result, ProviderActionQueueError> { + Self::provider_action_ids_to_enqueue_from( + &self.provider_tool_executions, + actions, + conversation_id, + batch, + ) + } + + fn provider_action_ids_to_enqueue_from( + existing_correlations: &HashMap< + (AIConversationId, AIAgentActionId), + ProviderToolExecutionRef, + >, + actions: &[AIAgentAction], + conversation_id: AIConversationId, + batch: &PendingToolBatch, + ) -> Result, ProviderActionQueueError> { + let refs = provider_action_correlations(actions, conversation_id, batch)?; + let mut action_ids = HashSet::with_capacity(refs.len()); + let mut seen_refs = HashSet::with_capacity(refs.len()); + for ((_, action_id), execution_ref) in &refs { + if !seen_refs.insert(execution_ref.clone()) { + return Err(ProviderActionQueueError::ExistingCorrelation { + call_id: action_id.to_string(), + }); + } + if let Some(existing_ref) = + existing_correlations.get(&(conversation_id, action_id.clone())) + { + if existing_ref == execution_ref { + continue; + } + return Err(ProviderActionQueueError::ExistingCorrelation { + call_id: action_id.to_string(), + }); + } + if existing_correlations + .values() + .any(|existing_ref| existing_ref == execution_ref) { return Err(ProviderActionQueueError::ExistingCorrelation { call_id: action_id.to_string(), }); } + action_ids.insert(action_id.clone()); } - self.provider_tool_executions.extend(refs); - self.executor.update(ctx, |executor, ctx| { - executor.mark_restored_actions(conversation_id, &recovery_action_ids, ctx); - }); - self.queue_actions(actions, conversation_id, ctx); - Ok(()) + Ok(action_ids) } /// Queues the `actions` in the given iterator for the given conversation, diff --git a/app/src/ai/blocklist/action_model_tests.rs b/app/src/ai/blocklist/action_model_tests.rs index fcbe8a35..7f47184d 100644 --- a/app/src/ai/blocklist/action_model_tests.rs +++ b/app/src/ai/blocklist/action_model_tests.rs @@ -104,6 +104,26 @@ fn provider_action_correlations_require_the_exact_unresolved_batch_order() { ); } +#[test] +fn exact_provider_action_correlation_is_idempotent() { + let conversation_id = AIConversationId::new(); + let batch = pending_tool_batch(&["first"]); + let actions = vec![action("first")]; + let execution_ref = ProviderToolExecutionRef::new(conversation_id, &batch.work_id, "first"); + let existing = HashMap::from([((conversation_id, actions[0].id.clone()), execution_ref)]); + + assert!( + BlocklistAIActionModel::provider_action_ids_to_enqueue_from( + &existing, + &actions, + conversation_id, + &batch, + ) + .unwrap() + .is_empty() + ); +} + #[test] fn parallel_phase_only_admits_matching_autoexecutable_actions() { let phase = diff --git a/app/src/ai/blocklist/block.rs b/app/src/ai/blocklist/block.rs index 6ab7c799..5d9db9b2 100644 --- a/app/src/ai/blocklist/block.rs +++ b/app/src/ai/blocklist/block.rs @@ -4390,7 +4390,9 @@ impl AIBlock { // If auto-login is enabled, run the login command automatically if auto_login_enabled { - ctx.emit(AIBlockEvent::RunAwsLoginCommand); + // Try the SDK chain first. This handles credentials that were refreshed by + // another AWS process without unnecessarily launching an interactive SSO flow. + ctx.emit(AIBlockEvent::RefreshAwsCredentials); } let model_name = model_name.clone(); @@ -6498,6 +6500,8 @@ pub enum AIBlockEvent { OpenActiveAgentProfileEditor, /// Run the configured AWS auth refresh command to fix expired Bedrock credentials RunAwsLoginCommand, + /// Reload credentials and fall back to the configured auth command if needed. + RefreshAwsCredentials, /// Emitted when a passive code diff has loaded its diffs and is ready to display. /// This is used to trigger height recalculation since the diffs are loaded asynchronously /// after the initial output completes. diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 1cb6032f..25109a53 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -5895,6 +5895,7 @@ impl BlocklistAIController { return; } }; + coordinator.set_max_context_tokens(response_config.max_context_tokens); if let Some(reason) = cancellation_reason { if !coordinator.run().is_terminal() { if let Err(error) = coordinator.run_mut().cancel(reason.to_string()) { @@ -6368,6 +6369,7 @@ impl BlocklistAIController { return; } }; + coordinator.set_max_context_tokens(response_config.max_context_tokens); if let Some(profile) = cli_monitor_profile { if let Err(error) = coordinator.insert_profile( CLI_MONITOR_PROVIDER_PROFILE, @@ -6961,7 +6963,7 @@ impl BlocklistAIController { ctx: &mut ModelContext, ) { let usage_u32 = |tokens: u64| u32::try_from(tokens).unwrap_or(u32::MAX); - let cost_cents = crate::ai::bedrock::response_translator::estimate_cost_cents( + let cost_cents = crate::ai::provider::response_translator::estimate_cost_cents( usage_u32(usage.input_tokens), usage_u32(usage.output_tokens), usage_u32(usage.cached_input_tokens), @@ -6997,6 +6999,13 @@ impl BlocklistAIController { return; } self.begin_active_provider_progressive_summary(conversation_id, ctx); + if matches!( + self.active_provider_progressive_summaries + .get(&conversation_id), + Some(ActiveProviderProgressiveSummaryState::InFlight { .. }) + ) { + return; + } let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { return; }; @@ -7452,6 +7461,24 @@ impl BlocklistAIController { } }; match block { + ProviderRunBlock::ContextWindowExceeded => { + slot.run = Some(run); + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + self.fail_active_provider_run( + conversation_id, + format!("failed to persist context compaction boundary: {error}"), + ctx, + ); + return; + } + if !self.begin_active_provider_progressive_summary(conversation_id, ctx) { + self.fail_active_provider_run( + conversation_id, + "provider context overflow could not start compaction".to_string(), + ctx, + ); + } + } ProviderRunBlock::ReadyToCallModel => { slot.run = Some(run); if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { @@ -7539,7 +7566,7 @@ impl BlocklistAIController { else { return; }; - let (converted_actions, invalid_results) = + let (mut converted_actions, invalid_results) = convert_provider_tool_batch(&run.action_context, &batch); for result in &invalid_results { if let Err(error) = run @@ -7586,6 +7613,30 @@ impl BlocklistAIController { }; } } + let action_ids_to_enqueue = match self.action_model.update(ctx, |action_model, _| { + action_model.provider_action_ids_to_enqueue( + &converted_actions + .iter() + .map(|(action, _)| action.clone()) + .collect::>(), + conversation_id, + &executable_batch, + ) + }) { + Ok(action_ids) => action_ids, + Err(error) => { + self.fail_active_provider_run(conversation_id, error.to_string(), ctx); + return; + } + }; + converted_actions.retain(|(action, _)| action_ids_to_enqueue.contains(&action.id)); + let queued_call_ids = converted_actions + .iter() + .map(|(action, _)| action.id.to_string()) + .collect::>(); + executable_batch + .calls + .retain(|call| queued_call_ids.contains(&call.call.id)); let stream_id = self.active_provider_runs[&conversation_id] .stream_id .clone(); @@ -9522,7 +9573,7 @@ impl BlocklistAIController { } Some(warp_multi_agent_api::response_event::stream_finished::Reason::ContextWindowExceeded(_)) => { let error_message = "Input exceeded context window limit."; - crate::ai::bedrock::crash_log::log_crash( + crate::ai::provider::crash_log::log_crash( "ContextWindowExceeded", error_message, "unknown", @@ -9615,7 +9666,7 @@ impl BlocklistAIController { let error_message = format!( "Response stream finished unexpectedly with internal error: {message}", ); - crate::ai::bedrock::crash_log::log_crash( + crate::ai::provider::crash_log::log_crash( "InternalError", &error_message, "unknown", @@ -9640,7 +9691,7 @@ impl BlocklistAIController { } Some(warp_multi_agent_api::response_event::stream_finished::Reason::MaxTokenLimit(_)) => { let error_message = "Input exceeded context window limit."; - crate::ai::bedrock::crash_log::log_crash( + crate::ai::provider::crash_log::log_crash( "MaxTokenLimit", error_message, "unknown", @@ -9940,7 +9991,7 @@ impl BlocklistAIController { u32::try_from(tokens).unwrap_or(u32::MAX) }; let cost_cents = - crate::ai::bedrock::response_translator::estimate_cost_cents( + crate::ai::provider::response_translator::estimate_cost_cents( usage_u32(usage.input_tokens), usage_u32(usage.output_tokens), usage_u32(usage.cached_input_tokens), diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index a03c5795..fdeb0903 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -32,11 +32,11 @@ use crate::ai::agent::conversation::AIConversationId; #[cfg(not(target_family = "wasm"))] use crate::ai::agent::AIAgentInput; use crate::ai::agent::{AIIdentifiers, CancellationReason}; -use crate::ai::bedrock::client::BedrockClientConfig; #[cfg(not(target_family = "wasm"))] use crate::ai::blocklist::BlocklistAIPermissions; use crate::ai::llms::{LLMId, LLMPreferences}; use crate::ai::openai::client::OpenAIClientConfig; +use crate::ai::provider::client::BedrockClientConfig; use crate::ai::provider::ProviderConfig; #[cfg(not(target_family = "wasm"))] use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord}; @@ -181,7 +181,6 @@ impl ResponseStream { reasoning_effort: client_config.reasoning_effort.clone(), max_input_tokens: client_config.max_input_tokens, max_output_tokens: client_config.max_output_tokens, - use_rig: client_config.use_rig, supports_system_messages: client_config.supports_system_messages, }); } @@ -191,12 +190,6 @@ impl ResponseStream { let auth_method = *settings.bedrock_auth_method.value(); let region = settings.bedrock_region.value().clone(); let cross_region_inference = *settings.bedrock_cross_region_inference.value(); - let use_rig = crate::ai::bedrock::models::configured_model_uses_rig( - model_id, - settings.bedrock_models.value(), - ®ion, - cross_region_inference, - ); let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx); let mut config = BedrockClientConfig { auth_method, @@ -206,7 +199,6 @@ impl ResponseStream { secret_access_key: settings.bedrock_secret_access_key.value().clone(), session_token: None, cross_region_inference, - use_rig, }; if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } = @@ -236,7 +228,7 @@ impl ResponseStream { format!("bedrock:{:?}:region={region}", config.auth_method) } ProviderConfig::OpenAI(config) => { - format!("openai:{:?}:rig={}", config.kind, config.use_rig) + format!("openai:{:?}", config.kind) } ProviderConfig::None => "none".to_string(), } diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index 9637cc7b..0e456ac9 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -267,6 +267,7 @@ fn provider_snapshot(conversation_id: AIConversationId) -> super::ActiveProvider max_context_tokens: Some(128_000), capabilities: RuntimeCapabilities::provider(), empty_output_message: None, + todo_items: None, }, action_context: crate::ai::runtime::ProviderActionContext::new_for_test( task_id.to_string(), diff --git a/app/src/ai/blocklist/usage/context_window_view.rs b/app/src/ai/blocklist/usage/context_window_view.rs index c28e042d..ced61776 100644 --- a/app/src/ai/blocklist/usage/context_window_view.rs +++ b/app/src/ai/blocklist/usage/context_window_view.rs @@ -3,7 +3,7 @@ use galaxyui::elements::{ }; use galaxyui::{AppContext, Element, Entity, SingletonEntity, TypedActionView, View, ViewContext}; -use crate::ai::bedrock::convert::{ContentPart, ConversationMessage, MessageContent, MessageRole}; +use crate::ai::provider::convert::{ContentPart, ConversationMessage, MessageContent, MessageRole}; use crate::appearance::Appearance; use crate::ui_components::blended_colors; diff --git a/app/src/ai/crosscheck/reviewer.rs b/app/src/ai/crosscheck/reviewer.rs index 98ff169e..2e6bbcca 100644 --- a/app/src/ai/crosscheck/reviewer.rs +++ b/app/src/ai/crosscheck/reviewer.rs @@ -175,7 +175,6 @@ impl CrosscheckReviewer { reasoning_effort: client_config.reasoning_effort.clone(), max_input_tokens: client_config.max_input_tokens, max_output_tokens: Some(REVIEWER_MAX_OUTPUT_TOKENS), - use_rig: client_config.use_rig, supports_system_messages: client_config.supports_system_messages, }); } @@ -184,7 +183,7 @@ impl CrosscheckReviewer { if *settings.bedrock_enabled.value() { let auth_method = *settings.bedrock_auth_method.value(); let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx); - let mut config = crate::ai::bedrock::client::BedrockClientConfig { + let mut config = crate::ai::provider::client::BedrockClientConfig { auth_method, profile: settings.bedrock_profile.value().clone(), region: settings.bedrock_region.value().clone(), @@ -192,7 +191,6 @@ impl CrosscheckReviewer { secret_access_key: settings.bedrock_secret_access_key.value().clone(), session_token: None, cross_region_inference: *settings.bedrock_cross_region_inference.value(), - use_rig: false, }; if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } = @@ -339,9 +337,9 @@ impl CrosscheckReviewer { async fn invoke_via_bedrock( agent_output: String, model_id: String, - config: crate::ai::bedrock::client::BedrockClientConfig, + config: crate::ai::provider::client::BedrockClientConfig, ) -> Result { - use crate::ai::bedrock::client::BedrockClient; + use crate::ai::provider::client::BedrockClient; use crate::ai::provider::types::{ConversationMessage, MessageContent, MessageRole}; let cross_region_inference = config.cross_region_inference; diff --git a/app/src/ai/llms.rs b/app/src/ai/llms.rs index ce0993ac..a5415298 100644 --- a/app/src/ai/llms.rs +++ b/app/src/ai/llms.rs @@ -779,7 +779,7 @@ impl LLMPreferences { if !*settings.bedrock_enabled.value() { return; } - let config = crate::ai::bedrock::client::BedrockClientConfig { + let config = crate::ai::provider::client::BedrockClientConfig { auth_method: *settings.bedrock_auth_method.value(), profile: settings.bedrock_profile.value().clone(), region: settings.bedrock_region.value().clone(), @@ -787,11 +787,10 @@ impl LLMPreferences { secret_access_key: settings.bedrock_secret_access_key.value().clone(), session_token: None, cross_region_inference: *settings.bedrock_cross_region_inference.value(), - use_rig: false, }; let _ = ctx.spawn( - async move { crate::ai::bedrock::discovery::discover_available_models(config).await }, + async move { crate::ai::provider::discovery::discover_available_models(config).await }, |me, result, ctx| match result { Ok(models) => { AISettings::handle(ctx).update(ctx, |settings, ctx| { @@ -839,7 +838,7 @@ impl LLMPreferences { let cross_region = *settings.bedrock_cross_region_inference.value(); // Check if user wants only 1-hour cache models - use crate::ai::bedrock::external_config::ExternalBedrockConfig; + use crate::ai::provider::external_config::ExternalBedrockConfig; let external_config = ExternalBedrockConfig::load(); let require_1h_cache = external_config.enable_prompt_caching_1h; @@ -885,7 +884,7 @@ impl LLMPreferences { } let model_id = if cross_region && !region.is_empty() { - super::bedrock::models::apply_cross_region_prefix(&model.model_id, ®ion) + super::provider::models::apply_cross_region_prefix(&model.model_id, ®ion) } else { model.model_id.clone() }; @@ -1173,11 +1172,6 @@ impl LLMPreferences { reasoning_effort: reasoning_effort.clone(), max_input_tokens: Some(openai_model_context_size(model)), max_output_tokens: model.max_output_tokens, - use_rig: model.use_rig - || !matches!( - provider_kind, - OpenAIProviderKind::OpenAI | OpenAIProviderKind::LiteLLM - ), supports_system_messages: model.supports_system_messages(), }; self.openai_provider_routing @@ -1841,7 +1835,6 @@ impl LLMPreferences { max_input_tokens: model.context_size, max_output_tokens: None, provider: None, - use_rig: true, supports_system_messages: Some(true), capability_overrides: std::collections::HashMap::new(), reasoning_efforts: Vec::new(), @@ -2658,15 +2651,12 @@ pub(crate) fn merge_discovered_provider_models( { discovered.display_name = existing.display_name.clone(); discovered.enabled = existing.enabled; - discovered.use_rig = existing.use_rig; if existing.supports_system_messages.is_some() { discovered.supports_system_messages = existing.supports_system_messages; } if discovered.provider.is_none() { discovered.provider = existing.provider.clone(); } - } else { - discovered.use_rig = true; } if discovered.model_id.starts_with("codex-gpt-") { discovered.supports_system_messages = Some(false); @@ -2707,7 +2697,6 @@ pub(crate) fn merge_discovered_chatgpt_subscription_models( .find(|model| model.model_id == discovered.model_id) { discovered.enabled = existing.enabled; - discovered.use_rig = existing.use_rig; if existing.supports_system_messages.is_some() { discovered.supports_system_messages = existing.supports_system_messages; } @@ -2934,7 +2923,6 @@ fn chatgpt_models_from_codex_response(body: &serde_json::Value) -> Vec OpenAIModelConfig { max_input_tokens: None, max_output_tokens: None, provider: None, - use_rig: false, supports_system_messages: None, capability_overrides: std::collections::HashMap::new(), reasoning_efforts: Vec::new(), @@ -164,7 +163,6 @@ fn bedrock_model(model_id: &str) -> BedrockModelConfig { model_id: model_id.to_string(), display_name: model_id.to_string(), vision_supported: false, - use_rig: true, } } @@ -308,7 +306,6 @@ fn provider_discovery_preserves_local_model_overrides() { existing.display_name = "My Codex".to_string(); existing.context_size = 100_000; existing.provider = Some("openai".to_string()); - existing.use_rig = true; // Even stale or incorrect endpoint metadata must not opt ChatGPT-backed // Codex models back into the system role. existing.supports_system_messages = Some(true); @@ -325,7 +322,6 @@ fn provider_discovery_preserves_local_model_overrides() { assert_eq!(merged[0].display_name, "My Codex"); assert_eq!(merged[0].context_size, 400_000); assert_eq!(merged[0].max_output_tokens, Some(32_000)); - assert!(merged[0].use_rig); assert_eq!(merged[0].supports_system_messages, Some(false)); assert_eq!(merged[0].provider.as_deref(), Some("openai")); } @@ -725,7 +721,7 @@ fn disabled_providers_do_not_leave_models_in_the_runtime_inventory() { } #[test] -fn provider_discovery_enables_rig_for_new_models_and_keeps_manual_models() { +fn provider_discovery_keeps_new_and_manual_models() { let manual = openai_model("manual-model"); let discovered = openai_model("codex-gpt-new"); @@ -733,7 +729,6 @@ fn provider_discovery_enables_rig_for_new_models_and_keeps_manual_models() { assert_eq!(merged.len(), 2); assert_eq!(merged[0].model_id, "codex-gpt-new"); - assert!(merged[0].use_rig); assert_eq!(merged[0].supports_system_messages, Some(false)); assert_eq!(merged[1].model_id, "manual-model"); } @@ -798,7 +793,6 @@ fn chatgpt_codex_models_parse_visible_catalog_entries() { assert_eq!(models[0].context_size, 828_400); assert_eq!(models[0].max_input_tokens, Some(258_400)); assert_eq!(models[0].reasoning_efforts, ["low", "xhigh", "ultra"]); - assert!(models[0].use_rig); assert_eq!(models[0].provider.as_deref(), Some("openai")); assert_eq!(models[0].supports_system_messages, Some(true)); @@ -828,7 +822,6 @@ fn chatgpt_codex_models_parse_visible_catalog_entries() { fn chatgpt_catalog_merge_drops_stale_models_but_preserves_overrides() { let mut existing = openai_model("gpt-5.6-sol"); existing.enabled = false; - existing.use_rig = false; existing.context_size = 272_000; existing.max_input_tokens = Some(272_000); existing.capability_overrides.insert( @@ -840,7 +833,6 @@ fn chatgpt_catalog_merge_drops_stale_models_but_preserves_overrides() { let mut discovered = openai_model("gpt-5.6-sol"); discovered.display_name = "GPT-5.6-Sol".to_string(); discovered.vision_supported = true; - discovered.use_rig = true; discovered.context_size = 828_400; discovered.max_input_tokens = Some(258_400); @@ -851,7 +843,6 @@ fn chatgpt_catalog_merge_drops_stale_models_but_preserves_overrides() { assert_eq!(merged[0].context_size, 828_400); assert_eq!(merged[0].max_input_tokens, Some(258_400)); assert!(!merged[0].enabled); - assert!(!merged[0].use_rig); assert_eq!( merged[0].capability_override("vision"), crate::settings::ModelCapabilityOverride::Unsupported diff --git a/app/src/ai/mod.rs b/app/src/ai/mod.rs index 01b7b781..20fc1b3a 100644 --- a/app/src/ai/mod.rs +++ b/app/src/ai/mod.rs @@ -19,8 +19,6 @@ pub mod auth_secret_types; #[cfg(not(target_family = "wasm"))] pub mod aws_credentials; #[cfg(not(target_family = "wasm"))] -pub mod bedrock; -#[cfg(not(target_family = "wasm"))] pub(crate) mod bedrock_credentials; pub(crate) mod block_context; pub(crate) mod blocklist; diff --git a/app/src/ai/openai/client.rs b/app/src/ai/openai/client.rs index 80771a54..23bb8eff 100644 --- a/app/src/ai/openai/client.rs +++ b/app/src/ai/openai/client.rs @@ -17,7 +17,6 @@ pub struct OpenAIClientConfig { pub reasoning_effort: Option, pub max_input_tokens: Option, pub max_output_tokens: Option, - pub use_rig: bool, pub supports_system_messages: bool, } diff --git a/app/src/ai/openai/response_translator.rs b/app/src/ai/openai/response_translator.rs index 4a36ca66..f074956b 100644 --- a/app/src/ai/openai/response_translator.rs +++ b/app/src/ai/openai/response_translator.rs @@ -10,7 +10,7 @@ use warp_multi_agent_api::response_event::stream_finished; use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent}; use crate::ai::agent::api::LegacyEvent; -use crate::ai::bedrock::response_translator::{ +use crate::ai::provider::response_translator::{ build_create_task, build_stream_init, context_window_for_model, }; use crate::ai::provider::types::{ContentPart, ConversationMessage, MessageContent, MessageRole}; @@ -531,7 +531,7 @@ fn build_tool_call_message( tool_input_json: &str, ) -> ResponseEvent { // Reuse the Bedrock tool call message builder since the proto output is identical - crate::ai::bedrock::response_translator::build_tool_call_message( + crate::ai::provider::response_translator::build_tool_call_message( task_id, tool_use_id, tool_name, diff --git a/app/src/ai/prompt_builder/mod.rs b/app/src/ai/prompt_builder/mod.rs index d4c1584c..53965c83 100644 --- a/app/src/ai/prompt_builder/mod.rs +++ b/app/src/ai/prompt_builder/mod.rs @@ -26,7 +26,7 @@ pub use prompts::provider::Provider; #[allow(unused_imports)] pub use tools::tools_for_mode; -use crate::ai::bedrock::convert::ToolDefinition; +use crate::ai::provider::convert::ToolDefinition; /// The fully-resolved prompt configuration ready to send to a model. #[derive(Debug, Clone)] diff --git a/app/src/ai/prompt_builder/tests.rs b/app/src/ai/prompt_builder/tests.rs index 801f0c44..c9f0466c 100644 --- a/app/src/ai/prompt_builder/tests.rs +++ b/app/src/ai/prompt_builder/tests.rs @@ -103,7 +103,7 @@ mod prompt_builder_tests { #[test] fn test_mcp_tools_appended() { - use crate::ai::bedrock::convert::ToolDefinition; + use crate::ai::provider::convert::ToolDefinition; let mcp_tool = ToolDefinition { name: "mcp__github__create_pr".to_string(), diff --git a/app/src/ai/prompt_builder/tools.rs b/app/src/ai/prompt_builder/tools.rs index 400aeea3..d9776e12 100644 --- a/app/src/ai/prompt_builder/tools.rs +++ b/app/src/ai/prompt_builder/tools.rs @@ -3,7 +3,7 @@ //! Each mode has a different set of tools available. Code mode gets everything, //! Plan mode gets read-only tools, Review mode gets read + search, etc. -use crate::ai::bedrock::convert::ToolDefinition; +use crate::ai::provider::convert::ToolDefinition; use crate::ai::prompt_builder::mode::Mode; /// Returns the tool definitions available for the given mode. diff --git a/app/src/ai/bedrock/client.rs b/app/src/ai/provider/client.rs similarity index 99% rename from app/src/ai/bedrock/client.rs rename to app/src/ai/provider/client.rs index a114c791..b77ce134 100644 --- a/app/src/ai/bedrock/client.rs +++ b/app/src/ai/provider/client.rs @@ -38,7 +38,6 @@ pub struct BedrockClientConfig { pub secret_access_key: String, pub session_token: Option, pub cross_region_inference: bool, - pub use_rig: bool, } impl BedrockClientConfig { diff --git a/app/src/ai/bedrock/convert.rs b/app/src/ai/provider/convert.rs similarity index 100% rename from app/src/ai/bedrock/convert.rs rename to app/src/ai/provider/convert.rs diff --git a/app/src/ai/bedrock/convert_request.rs b/app/src/ai/provider/convert_request.rs similarity index 100% rename from app/src/ai/bedrock/convert_request.rs rename to app/src/ai/provider/convert_request.rs diff --git a/app/src/ai/bedrock/convert_tests.rs b/app/src/ai/provider/convert_tests.rs similarity index 100% rename from app/src/ai/bedrock/convert_tests.rs rename to app/src/ai/provider/convert_tests.rs diff --git a/app/src/ai/bedrock/crash_log.rs b/app/src/ai/provider/crash_log.rs similarity index 100% rename from app/src/ai/bedrock/crash_log.rs rename to app/src/ai/provider/crash_log.rs diff --git a/app/src/ai/bedrock/diagnostic.rs b/app/src/ai/provider/diagnostic.rs similarity index 100% rename from app/src/ai/bedrock/diagnostic.rs rename to app/src/ai/provider/diagnostic.rs diff --git a/app/src/ai/bedrock/discovery.rs b/app/src/ai/provider/discovery.rs similarity index 99% rename from app/src/ai/bedrock/discovery.rs rename to app/src/ai/provider/discovery.rs index d9e673d5..1e40f839 100644 --- a/app/src/ai/bedrock/discovery.rs +++ b/app/src/ai/provider/discovery.rs @@ -88,7 +88,6 @@ pub async fn discover_available_models( model_id: model_id.to_owned(), display_name, vision_supported, - use_rig: true, }) } }); diff --git a/app/src/ai/bedrock/e2e_tests.rs b/app/src/ai/provider/e2e_tests.rs similarity index 99% rename from app/src/ai/bedrock/e2e_tests.rs rename to app/src/ai/provider/e2e_tests.rs index 3cfa7111..aedb3563 100644 --- a/app/src/ai/bedrock/e2e_tests.rs +++ b/app/src/ai/provider/e2e_tests.rs @@ -269,7 +269,6 @@ fn get_test_config() -> Option { secret_access_key: String::new(), session_token: None, cross_region_inference: false, - use_rig: false, }) } @@ -280,7 +279,7 @@ fn get_test_model() -> String { } fn sample_project_path() -> PathBuf { - PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("src/ai/bedrock/test_fixtures/sample_project") + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("src/ai/provider/test_fixtures/sample_project") } fn agent_tools() -> Vec { diff --git a/app/src/ai/bedrock/external_config.rs b/app/src/ai/provider/external_config.rs similarity index 99% rename from app/src/ai/bedrock/external_config.rs rename to app/src/ai/provider/external_config.rs index 2f9cd94c..3cff89b6 100644 --- a/app/src/ai/bedrock/external_config.rs +++ b/app/src/ai/provider/external_config.rs @@ -144,7 +144,6 @@ fn parse_claude_code_model_map( model_id: arn, display_name, vision_supported: true, - use_rig: false, } }) .collect() diff --git a/app/src/ai/bedrock/external_config_tests.rs b/app/src/ai/provider/external_config_tests.rs similarity index 100% rename from app/src/ai/bedrock/external_config_tests.rs rename to app/src/ai/provider/external_config_tests.rs diff --git a/app/src/ai/bedrock/integration_tests.rs b/app/src/ai/provider/integration_tests.rs similarity index 99% rename from app/src/ai/bedrock/integration_tests.rs rename to app/src/ai/provider/integration_tests.rs index 27ceb715..155b3424 100644 --- a/app/src/ai/bedrock/integration_tests.rs +++ b/app/src/ai/provider/integration_tests.rs @@ -24,7 +24,6 @@ fn get_test_config() -> Option { secret_access_key: String::new(), session_token: None, cross_region_inference: false, - use_rig: false, }) } diff --git a/app/src/ai/provider/mod.rs b/app/src/ai/provider/mod.rs index e67d252f..498d84a5 100644 --- a/app/src/ai/provider/mod.rs +++ b/app/src/ai/provider/mod.rs @@ -1,12 +1,56 @@ pub mod types; -use crate::ai::bedrock::client::BedrockClientConfig; use crate::ai::openai::client::OpenAIClientConfig; +#[cfg(not(target_family = "wasm"))] +use crate::ai::provider::client::BedrockClientConfig; + #[allow(dead_code)] #[derive(Clone)] pub enum ProviderConfig { + #[cfg(not(target_family = "wasm"))] Bedrock(BedrockClientConfig), OpenAI(OpenAIClientConfig), None, } + +// Provider-neutral module boundary; Bedrock-specific transport helpers live here without a +// provider-named directory so the provider layer can be reorganized independently of callers. +#[cfg(not(target_family = "wasm"))] +pub mod client; +#[cfg(not(target_family = "wasm"))] +pub mod convert; +#[cfg(not(target_family = "wasm"))] +pub mod crash_log; +#[cfg(not(target_family = "wasm"))] +pub mod diagnostic; +#[cfg(not(target_family = "wasm"))] +pub mod discovery; +#[cfg(not(target_family = "wasm"))] +pub mod external_config; +#[cfg(not(target_family = "wasm"))] +pub mod models; +#[cfg(not(target_family = "wasm"))] +pub mod request_translator; +#[cfg(not(target_family = "wasm"))] +pub mod response_translator; +#[cfg(not(target_family = "wasm"))] +pub mod settings_view; + +#[cfg(not(target_family = "wasm"))] +#[cfg(test)] +mod convert_tests; +#[cfg(not(target_family = "wasm"))] +#[cfg(test)] +#[allow(dead_code)] +mod e2e_tests; +#[cfg(not(target_family = "wasm"))] +#[cfg(test)] +#[allow(dead_code)] +mod integration_tests; +#[cfg(not(target_family = "wasm"))] +#[cfg(test)] +mod models_tests; +#[cfg(not(target_family = "wasm"))] +#[cfg(test)] +mod response_translator_tests; diff --git a/app/src/ai/provider/models.rs b/app/src/ai/provider/models.rs new file mode 100644 index 00000000..9c9c69be --- /dev/null +++ b/app/src/ai/provider/models.rs @@ -0,0 +1,19 @@ +pub fn apply_cross_region_prefix(model_id: &str, region: &str) -> String { + if model_id.starts_with("arn:") { + return model_id.to_string(); + } + + if model_id.contains('.') && model_id.split('.').next().unwrap_or("").len() <= 6 { + return model_id.to_string(); + } + + let prefix = match region { + r if r.starts_with("us-") || r.starts_with("ca-") => "us", + r if r.starts_with("eu-") || r == "il-central-1" => "eu", + r if r == "ap-northeast-1" || r == "ap-northeast-3" => "jp", + r if r == "ap-southeast-2" || r == "ap-southeast-4" || r == "ap-southeast-6" => "au", + r if r.starts_with("ap-") => "apac", + _ => return model_id.to_string(), + }; + format!("{}.{}", prefix, model_id) +} diff --git a/app/src/ai/bedrock/models_tests.rs b/app/src/ai/provider/models_tests.rs similarity index 74% rename from app/src/ai/bedrock/models_tests.rs rename to app/src/ai/provider/models_tests.rs index 55f927ae..0331c740 100644 --- a/app/src/ai/bedrock/models_tests.rs +++ b/app/src/ai/provider/models_tests.rs @@ -1,6 +1,4 @@ use super::models::*; -use crate::settings::ai::BedrockModelConfig; - #[test] fn test_cross_region_prefix_us_east() { assert_eq!( @@ -83,32 +81,3 @@ fn test_cross_region_prefix_skips_arn() { let arn = "arn:aws:bedrock:us-east-1:156729649053:application-inference-profile/fma5bsw4oxhy"; assert_eq!(apply_cross_region_prefix(arn, "us-east-1"), arn); } - -#[test] -fn rig_opt_in_matches_context_markers_and_resolved_inference_profiles() { - let configured = vec![BedrockModelConfig { - model_id: "anthropic.claude-test[1m]".to_string(), - display_name: "Claude Test".to_string(), - vision_supported: false, - use_rig: true, - }]; - - assert!(configured_model_uses_rig( - "us.anthropic.claude-test", - &configured, - "us-east-1", - true, - )); - assert!(configured_model_uses_rig( - "anthropic.claude-test[1M]", - &configured, - "us-east-1", - false, - )); - assert!(!configured_model_uses_rig( - "anthropic.other-model", - &configured, - "us-east-1", - false, - )); -} diff --git a/app/src/ai/bedrock/request_translator.rs b/app/src/ai/provider/request_translator.rs similarity index 100% rename from app/src/ai/bedrock/request_translator.rs rename to app/src/ai/provider/request_translator.rs diff --git a/app/src/ai/bedrock/request_translator_tests.rs b/app/src/ai/provider/request_translator_tests.rs similarity index 99% rename from app/src/ai/bedrock/request_translator_tests.rs rename to app/src/ai/provider/request_translator_tests.rs index 4eb307a3..e3796ce4 100644 --- a/app/src/ai/bedrock/request_translator_tests.rs +++ b/app/src/ai/provider/request_translator_tests.rs @@ -6,7 +6,7 @@ use super::{ extract_tools, inject_input_messages_into_task, sanitize_messages_for_bedrock, }; use crate::ai::agent::api::is_internal_command_completion_assessment; -use crate::ai::bedrock::convert::{ContentPart, ConversationMessage, MessageContent, MessageRole}; +use crate::ai::provider::convert::{ContentPart, ConversationMessage, MessageContent, MessageRole}; #[test] fn test_sanitize_messages_prepends_synthetic_tool_result_before_existing_user_text() { diff --git a/app/src/ai/bedrock/response_translator.rs b/app/src/ai/provider/response_translator.rs similarity index 100% rename from app/src/ai/bedrock/response_translator.rs rename to app/src/ai/provider/response_translator.rs diff --git a/app/src/ai/bedrock/response_translator_tests.rs b/app/src/ai/provider/response_translator_tests.rs similarity index 100% rename from app/src/ai/bedrock/response_translator_tests.rs rename to app/src/ai/provider/response_translator_tests.rs diff --git a/app/src/ai/bedrock/settings_view.rs b/app/src/ai/provider/settings_view.rs similarity index 98% rename from app/src/ai/bedrock/settings_view.rs rename to app/src/ai/provider/settings_view.rs index e036c3b6..d7ebefd9 100644 --- a/app/src/ai/bedrock/settings_view.rs +++ b/app/src/ai/provider/settings_view.rs @@ -3,8 +3,8 @@ use galaxyui::elements::{ }; use galaxyui::{AppContext, Element, Entity, SingletonEntity, TypedActionView, View, ViewContext}; -use crate::ai::bedrock::convert::CachingConfig; -use crate::ai::bedrock::external_config::ExternalBedrockConfig; +use crate::ai::provider::convert::CachingConfig; +use crate::ai::provider::external_config::ExternalBedrockConfig; use crate::appearance::Appearance; use crate::ui_components::blended_colors; diff --git a/app/src/ai/bedrock/test_fixtures/sample_project/Cargo.toml b/app/src/ai/provider/test_fixtures/sample_project/Cargo.toml similarity index 100% rename from app/src/ai/bedrock/test_fixtures/sample_project/Cargo.toml rename to app/src/ai/provider/test_fixtures/sample_project/Cargo.toml diff --git a/app/src/ai/bedrock/test_fixtures/sample_project/README.md b/app/src/ai/provider/test_fixtures/sample_project/README.md similarity index 100% rename from app/src/ai/bedrock/test_fixtures/sample_project/README.md rename to app/src/ai/provider/test_fixtures/sample_project/README.md diff --git a/app/src/ai/bedrock/test_fixtures/sample_project/src/lib.rs b/app/src/ai/provider/test_fixtures/sample_project/src/lib.rs similarity index 100% rename from app/src/ai/bedrock/test_fixtures/sample_project/src/lib.rs rename to app/src/ai/provider/test_fixtures/sample_project/src/lib.rs diff --git a/app/src/ai/bedrock/test_fixtures/sample_project/src/lib_tests.rs b/app/src/ai/provider/test_fixtures/sample_project/src/lib_tests.rs similarity index 100% rename from app/src/ai/bedrock/test_fixtures/sample_project/src/lib_tests.rs rename to app/src/ai/provider/test_fixtures/sample_project/src/lib_tests.rs diff --git a/app/src/ai/bedrock/test_fixtures/sample_project/src/main.rs b/app/src/ai/provider/test_fixtures/sample_project/src/main.rs similarity index 100% rename from app/src/ai/bedrock/test_fixtures/sample_project/src/main.rs rename to app/src/ai/provider/test_fixtures/sample_project/src/main.rs diff --git a/app/src/ai/bedrock/test_fixtures/sample_project/src/utils.rs b/app/src/ai/provider/test_fixtures/sample_project/src/utils.rs similarity index 100% rename from app/src/ai/bedrock/test_fixtures/sample_project/src/utils.rs rename to app/src/ai/provider/test_fixtures/sample_project/src/utils.rs diff --git a/app/src/ai/runtime/event_translator.rs b/app/src/ai/runtime/event_translator.rs index cd328e22..69e40cbe 100644 --- a/app/src/ai/runtime/event_translator.rs +++ b/app/src/ai/runtime/event_translator.rs @@ -10,11 +10,11 @@ use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent}; use super::provider_run_coordinator::ProviderRunProjection; use crate::ai::agent::runtime_activity; -use crate::ai::bedrock::response_translator::{ +use crate::ai::openai::response_translator::{build_stream_finished, StreamUsage}; +use crate::ai::provider::response_translator::{ build_add_agent_output_message, build_append_text, build_create_task, build_stream_init, build_user_query_message, }; -use crate::ai::openai::response_translator::{build_stream_finished, StreamUsage}; #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] pub(crate) struct RuntimeResponseConfig { @@ -26,6 +26,9 @@ pub(crate) struct RuntimeResponseConfig { pub(crate) max_context_tokens: Option, pub(crate) capabilities: RuntimeCapabilities, pub(crate) empty_output_message: Option, + /// Todo items supplied by the existing task transcript, when one is available. + #[serde(skip)] + pub(crate) todo_items: Option>, } /// Converts the provider-neutral runtime lifecycle into Galaxy's existing @@ -51,6 +54,7 @@ pub(crate) struct RuntimeResponseTranslator { pub(crate) struct ProviderRunResponseProjector { translator: RuntimeResponseTranslator, has_started_model_turn: bool, + todo_phase: usize, finished: bool, } @@ -59,6 +63,7 @@ impl ProviderRunResponseProjector { Self { translator: RuntimeResponseTranslator::new(config), has_started_model_turn: false, + todo_phase: 0, finished: false, } } @@ -70,6 +75,8 @@ impl ProviderRunResponseProjector { Self { translator: RuntimeResponseTranslator::restored(config, projection_was_initialized), has_started_model_turn: false, + // Task-list events are part of the already persisted projection. + todo_phase: usize::MAX, finished: false, } } @@ -87,17 +94,46 @@ impl ProviderRunResponseProjector { self.translator.begin_followup_turn(); } self.has_started_model_turn = true; - self.translator.translate(AgentEvent::TurnStarted { + let mut events = self.translator.translate(AgentEvent::TurnStarted { runtime_request_id: String::new(), - }) + })?; + events.extend(self.todo_phase_events()); + Ok(events) } ProviderRunProjection::ModelEvent { event, .. } => self.translator.translate(event), ProviderRunProjection::ModelRetry { .. } => { Ok(self.translator.discard_failed_turn_output()) } ProviderRunProjection::ModelTurnRequested { .. } - | ProviderRunProjection::ModelTurnFinished { .. } - | ProviderRunProjection::ToolBatchReady { .. } => Ok(Vec::new()), + | ProviderRunProjection::ModelTurnFinished { .. } => Ok(Vec::new()), + ProviderRunProjection::ToolBatchReady { .. } => { + let todo_index = self.todo_phase.saturating_sub(1); + let Some(todo) = self.todo_items().get(todo_index).cloned() else { + return Ok(Vec::new()); + }; + self.todo_phase += 1; + let mut events = vec![build_todo_update( + &self.translator.config.task_id, + api::message::update_todos::Operation::MarkTodosCompleted( + api::MarkTodosCompleted { + todo_ids: vec![todo.id], + }, + ), + )]; + events.push(build_todo_update( + &self.translator.config.task_id, + api::message::update_todos::Operation::UpdatePendingTodos( + api::UpdatePendingTodos { + updated_pending_todos: self + .todo_items() + .into_iter() + .skip(self.todo_phase) + .collect(), + }, + ), + )); + Ok(events) + } } } @@ -115,17 +151,154 @@ impl ProviderRunResponseProjector { } self.finished = true; match outcome { - ProviderRunOutcome::Completed(completion) => Ok(self - .translator - .finish_provider_run(completion.stop_reason.clone(), aggregate_usage)), - ProviderRunOutcome::Failed(failure) => Ok(self - .translator - .provider_failure(&failure.message, aggregate_usage)), + ProviderRunOutcome::Completed(completion) => { + let mut events = self.todo_completion_events(); + events.extend( + self.translator + .finish_provider_run(completion.stop_reason.clone(), aggregate_usage), + ); + Ok(events) + } + ProviderRunOutcome::Failed(failure) => Ok(self.translator.provider_failure( + &failure.message, + failure.source.as_ref(), + aggregate_usage, + )), ProviderRunOutcome::Cancelled { .. } => Ok(self .translator .finish_provider_run(StopReason::Cancelled, aggregate_usage)), } } + + // Keep the direct-provider workflow visible in the existing task list protocol. These are + // response events, so the normal history model remains the sole owner of task-list state. + fn todo_phase_events(&mut self) -> Vec { + let events = match self.todo_phase { + 0 => vec![build_todo_update( + &self.translator.config.task_id, + api::message::update_todos::Operation::CreateTodoList(api::CreateTodoList { + initial_todos: self.todo_items(), + }), + )], + _ => Vec::new(), + }; + // The first model turn owns the first phase. Tool batches advance it; this keeps + // arbitrary plan lengths aligned with the UpdateTodos protocol. + if self.todo_phase == 0 { + self.todo_phase = 1; + } + events + } + + fn todo_completion_events(&self) -> Vec { + if self.todo_phase == 0 || self.todo_phase == usize::MAX { + return Vec::new(); + } + let todos = self.todo_items(); + if todos.is_empty() { + return Vec::new(); + } + vec![ + build_todo_update( + &self.translator.config.task_id, + api::message::update_todos::Operation::MarkTodosCompleted( + api::MarkTodosCompleted { + todo_ids: todos.iter().map(|todo| todo.id.clone()).collect(), + }, + ), + ), + build_todo_update( + &self.translator.config.task_id, + api::message::update_todos::Operation::UpdatePendingTodos( + api::UpdatePendingTodos { + updated_pending_todos: Vec::new(), + }, + ), + ), + ] + } + fn todo_items(&self) -> Vec { + self.translator + .config + .todo_items + .clone() + .unwrap_or_else(default_workflow_todos) + } +} + +fn default_workflow_todos() -> Vec { + [ + api::TodoItem { + id: "direct-provider-research".to_owned(), + title: "Research the request".to_owned(), + description: "Inspect the repository and gather relevant evidence".to_owned(), + }, + api::TodoItem { + id: "direct-provider-plan".to_owned(), + title: "Create an implementation plan".to_owned(), + description: "Choose an approach grounded in the repository".to_owned(), + }, + api::TodoItem { + id: "direct-provider-critique".to_owned(), + title: "Critique the approach".to_owned(), + description: "Check assumptions, risks, and missing cases".to_owned(), + }, + api::TodoItem { + id: "direct-provider-revise".to_owned(), + title: "Revise the plan".to_owned(), + description: "Incorporate findings before editing".to_owned(), + }, + api::TodoItem { + id: "direct-provider-implement".to_owned(), + title: "Implement the change".to_owned(), + description: "Make the requested edits".to_owned(), + }, + api::TodoItem { + id: "direct-provider-verify".to_owned(), + title: "Verify the result".to_owned(), + description: "Run proportionate checks".to_owned(), + }, + api::TodoItem { + id: "direct-provider-repair".to_owned(), + title: "Repair validation issues".to_owned(), + description: "Fix failures found during verification".to_owned(), + }, + ] + .to_vec() +} + +fn build_todo_update( + task_id: &str, + operation: api::message::update_todos::Operation, +) -> ResponseEvent { + let message = api::Message { + id: Uuid::new_v4().to_string(), + task_id: task_id.to_owned(), + request_id: String::new(), + timestamp: None, + server_message_data: String::new(), + citations: Vec::new(), + fetched_memories: Vec::new(), + message: Some(api::message::Message::UpdateTodos( + api::message::UpdateTodos { + operation: Some(operation), + }, + )), + }; + ResponseEvent { + r#type: Some(api::response_event::Type::ClientActions( + api::response_event::ClientActions { + actions: vec![api::ClientAction { + action: Some(api::client_action::Action::AddMessagesToTask( + api::client_action::AddMessagesToTask { + task_id: task_id.to_owned(), + messages: vec![message], + }, + )), + }], + }, + )), + } } impl RuntimeResponseTranslator { @@ -421,15 +594,27 @@ impl RuntimeResponseTranslator { self.finished_with_reason(map_stop_reason(reason)) } - fn provider_failure(&mut self, message: &str, aggregate_usage: &Usage) -> Vec { + fn provider_failure( + &mut self, + message: &str, + source: Option<&galaxy_agent_core::AgentError>, + aggregate_usage: &Usage, + ) -> Vec { let mut events = Vec::new(); self.initialize(&mut events); - events.push(self.finished_with_usage( + let reason = if source + .is_some_and(|error| error.kind == galaxy_agent_core::AgentErrorKind::Authentication) + { + stream_finished::Reason::InvalidApiKey(stream_finished::InvalidApiKey { + provider: warp_multi_agent_api::LlmProvider::AwsBedrock as i32, + model_name: self.config.model_id.clone(), + }) + } else { stream_finished::Reason::InternalError(stream_finished::InternalError { message: message.to_owned(), - }), - aggregate_usage, - )); + }) + }; + events.push(self.finished_with_usage(reason, aggregate_usage)); events } diff --git a/app/src/ai/runtime/event_translator_tests.rs b/app/src/ai/runtime/event_translator_tests.rs index b3ad938e..0640ac61 100644 --- a/app/src/ai/runtime/event_translator_tests.rs +++ b/app/src/ai/runtime/event_translator_tests.rs @@ -17,6 +17,7 @@ fn provider_translator() -> RuntimeResponseTranslator { max_context_tokens: Some(1_000), capabilities: RuntimeCapabilities::provider(), empty_output_message: None, + todo_items: None, }) } @@ -30,6 +31,7 @@ fn session_translator() -> RuntimeResponseTranslator { max_context_tokens: None, capabilities: RuntimeCapabilities::session_runtime(), empty_output_message: Some("> runtime completed without text".to_owned()), + todo_items: None, }) } @@ -44,6 +46,7 @@ fn restored_provider_projection_skips_stream_initialization() { max_context_tokens: Some(1_000), capabilities: RuntimeCapabilities::provider(), empty_output_message: None, + todo_items: None, }; let mut projector = ProviderRunResponseProjector::restored(config, true); let work_id = galaxy_agent_core::ExternalWorkId { @@ -82,6 +85,109 @@ fn restored_provider_projection_skips_stream_initialization() { )); } +#[test] +fn provider_projection_uses_supplied_todos_and_advances_each_id() { + let todos = vec![ + warp_multi_agent_api::TodoItem { + id: "research".to_owned(), + title: "Research".to_owned(), + description: "Inspect".to_owned(), + }, + warp_multi_agent_api::TodoItem { + id: "implement".to_owned(), + title: "Implement".to_owned(), + description: "Edit".to_owned(), + }, + warp_multi_agent_api::TodoItem { + id: "verify".to_owned(), + title: "Verify".to_owned(), + description: "Check".to_owned(), + }, + ]; + let mut projector = ProviderRunResponseProjector::new(RuntimeResponseConfig { + task_id: "task".to_owned(), + conversation_id: "conversation".to_owned(), + needs_create_task: false, + user_query: None, + model_id: "model".to_owned(), + max_context_tokens: Some(1_000), + capabilities: RuntimeCapabilities::provider(), + empty_output_message: None, + todo_items: Some(todos), + }); + let work_id = galaxy_agent_core::ExternalWorkId { + run_id: galaxy_agent_core::ProviderRunId::new("run"), + epoch: galaxy_agent_core::RunEpoch::new(1), + }; + let initial = projector + .project(ProviderRunProjection::ModelTurnStarted { + work_id: work_id.clone(), + profile: galaxy_agent_core::ProviderRequestProfile::new("base"), + runtime_id: "runtime".to_owned(), + model_id: "model".to_owned(), + runtime_request_id: "request".to_owned(), + retry_attempt: 0, + elapsed_ms: 1, + }) + .unwrap(); + assert_eq!(initial.len(), 2); + let first = projector + .project(ProviderRunProjection::ToolBatchReady { + batch: galaxy_agent_core::PendingToolBatch { + work_id: galaxy_agent_core::ExternalWorkId { + run_id: galaxy_agent_core::ProviderRunId::new("run"), + epoch: galaxy_agent_core::RunEpoch::new(1), + }, + calls: Vec::new(), + }, + }) + .unwrap(); + let second = projector + .project(ProviderRunProjection::ToolBatchReady { + batch: galaxy_agent_core::PendingToolBatch { + work_id: galaxy_agent_core::ExternalWorkId { + run_id: galaxy_agent_core::ProviderRunId::new("run"), + epoch: galaxy_agent_core::RunEpoch::new(1), + }, + calls: Vec::new(), + }, + }) + .unwrap(); + let ids = |events: &[warp_multi_agent_api::ResponseEvent]| { + events + .iter() + .flat_map(|event| match &event.r#type { + Some(response_event::Type::ClientActions(actions)) => actions + .actions + .iter() + .filter_map(|action| match &action.action { + Some(client_action::Action::AddMessagesToTask(add)) => add + .messages + .iter() + .filter_map(|message| match &message.message { + Some(message::Message::UpdateTodos(update)) => match update + .operation + .as_ref() + { + Some(message::update_todos::Operation::MarkTodosCompleted( + mark, + )) => Some(mark.todo_ids[0].clone()), + _ => None, + }, + _ => None, + }) + .next(), + _ => None, + }) + .collect::>(), + _ => Vec::new(), + }) + .collect::>() + }; + assert_eq!(ids(&first), vec!["research"]); + assert_eq!(ids(&second), vec!["implement"]); +} + #[test] fn restored_uninitialized_projection_replays_init_before_live_delta() { let config = RuntimeResponseConfig { @@ -93,6 +199,7 @@ fn restored_uninitialized_projection_replays_init_before_live_delta() { max_context_tokens: Some(1_000), capabilities: RuntimeCapabilities::provider(), empty_output_message: None, + todo_items: None, }; let mut projector = ProviderRunResponseProjector::restored(config, false); let work_id = galaxy_agent_core::ExternalWorkId { @@ -120,11 +227,18 @@ fn restored_uninitialized_projection_replays_init_before_live_delta() { }) .unwrap(); - assert_eq!(started.len(), 1); + assert_eq!(started.len(), 2); assert!(matches!( started[0].r#type, Some(response_event::Type::Init(_)) )); + let Some(response_event::Type::ClientActions(actions)) = &started[1].r#type else { + panic!("initial provider turn should publish its task list"); + }; + assert!(matches!( + actions.actions[0].action, + Some(client_action::Action::AddMessagesToTask(_)) + )); assert_eq!(delta.len(), 1); assert!(matches!( delta[0].r#type, @@ -143,6 +257,7 @@ fn provider_followup_turn_starts_a_distinct_text_message() { max_context_tokens: Some(1_000), capabilities: RuntimeCapabilities::provider(), empty_output_message: None, + todo_items: None, }); let first_work_id = galaxy_agent_core::ExternalWorkId { run_id: galaxy_agent_core::ProviderRunId::new("run"), @@ -366,6 +481,7 @@ fn provider_retry_clears_failed_attempt_output_before_new_messages() { max_context_tokens: Some(1_000), capabilities: RuntimeCapabilities::provider(), empty_output_message: None, + todo_items: None, }); let work_id = galaxy_agent_core::ExternalWorkId { run_id: galaxy_agent_core::ProviderRunId::new("run"), diff --git a/app/src/ai/runtime/provider_run_coordinator.rs b/app/src/ai/runtime/provider_run_coordinator.rs index fa9479bd..91c430fa 100644 --- a/app/src/ai/runtime/provider_run_coordinator.rs +++ b/app/src/ai/runtime/provider_run_coordinator.rs @@ -82,6 +82,7 @@ pub(crate) enum ProviderRunProjection { pub(crate) enum ProviderRunBlock { Tools(PendingToolBatch), ReadyToCallModel, + ContextWindowExceeded, AwaitingDriver { work_id: ExternalWorkId, stop_reason: StopReason, @@ -175,6 +176,8 @@ pub(crate) struct ProviderRunCoordinator { profiles: BTreeMap, model_start_timeout: Duration, model_event_idle_timeout: Duration, + context_compaction_requested: bool, + max_context_tokens: Option, } impl ProviderRunCoordinator { @@ -218,6 +221,8 @@ impl ProviderRunCoordinator { profiles, model_start_timeout: PROVIDER_MODEL_START_TIMEOUT, model_event_idle_timeout: PROVIDER_MODEL_EVENT_IDLE_TIMEOUT, + context_compaction_requested: false, + max_context_tokens: None, }) } @@ -239,6 +244,10 @@ impl ProviderRunCoordinator { self.profiles.get(profile).map(|profile| &profile.request) } + pub(crate) fn set_max_context_tokens(&mut self, max_context_tokens: Option) { + self.max_context_tokens = max_context_tokens; + } + pub(crate) fn insert_profile( &mut self, profile: impl Into, @@ -364,8 +373,16 @@ impl ProviderRunCoordinator { if !self.checkpoint_or_fail(&mut checkpoint).await? { return self.terminal_block(); } + if self.request_needs_context_compaction(&call) { + self.run.prepare_context_compaction(&call.work_id)?; + return Ok(ProviderRunBlock::ContextWindowExceeded); + } self.drive_model_call_acknowledged(call, control.clone(), &mut project) .await?; + if self.context_compaction_requested { + self.context_compaction_requested = false; + return Ok(ProviderRunBlock::ContextWindowExceeded); + } } Some(ProviderRunStep::DispatchTools(batch)) => { if !self.checkpoint_or_fail(&mut checkpoint).await? { @@ -428,6 +445,17 @@ impl ProviderRunCoordinator { } } + fn request_needs_context_compaction(&self, call: &ProviderModelCall) -> bool { + let Some(max_context_tokens) = self.max_context_tokens else { + return false; + }; + let Some(profile) = self.profiles.get(call.profile.as_str()) else { + return false; + }; + let request = request_for_model_call(profile.request.clone(), call); + estimate_turn_request_tokens(&request) >= u64::from(max_context_tokens) + } + async fn checkpoint_or_fail( &mut self, checkpoint: &mut C, @@ -830,6 +858,11 @@ impl ProviderRunCoordinator { self.run.cancel("provider model call was cancelled")?; return Ok(()); } + if reason == StopReason::ContextWindowExceeded { + self.run.prepare_context_compaction(&call.work_id)?; + self.context_compaction_requested = true; + return Ok(()); + } let turn = buffer.complete(reason, advertised_tools); if let Err(error) = self.run.accept_model_turn(&call.work_id, turn) { self.run.fail( @@ -925,6 +958,11 @@ impl ProviderRunCoordinator { "payload": &error, })); let error_message = error.message.clone(); + if error.kind == AgentErrorKind::ContextWindowExceeded { + self.run.prepare_context_compaction(&call.work_id)?; + self.context_compaction_requested = true; + return Ok(()); + } let disposition = self .run .register_model_failure(&call.work_id, error.clone())?; @@ -1082,6 +1120,44 @@ fn request_for_model_call(mut template: TurnRequest, call: &ProviderModelCall) - template } +const ESTIMATED_CHARS_PER_TOKEN: u64 = 4; + +/// Deliberately overestimates request size without requiring provider-specific tokenizers. +fn estimate_turn_request_tokens(request: &TurnRequest) -> u64 { + let mut chars = request + .system_prompt + .as_deref() + .map_or(0, |text| text.chars().count()); + chars += request.prompt.as_ref().map_or(0, |prompt| { + serde_json::to_string(prompt) + .unwrap_or_default() + .chars() + .count() + }); + chars += request + .messages + .iter() + .map(|message| { + serde_json::to_string(message) + .unwrap_or_default() + .chars() + .count() + }) + .sum::(); + chars += request + .tools + .iter() + .map(|tool| { + serde_json::to_string(tool) + .unwrap_or_default() + .chars() + .count() + }) + .sum::(); + let input_tokens = (chars as u64).div_ceil(ESTIMATED_CHARS_PER_TOKEN); + input_tokens.saturating_add(request.max_output_tokens.unwrap_or_default()) +} + fn tool_event_call_id(event: &ToolEvent) -> Result<&str, ProviderRunCoordinatorError> { match event { ToolEvent::Proposed { call } => Ok(&call.id), diff --git a/app/src/ai/runtime/provider_run_coordinator_tests.rs b/app/src/ai/runtime/provider_run_coordinator_tests.rs index 8def1ec3..934d0c6a 100644 --- a/app/src/ai/runtime/provider_run_coordinator_tests.rs +++ b/app/src/ai/runtime/provider_run_coordinator_tests.rs @@ -164,6 +164,48 @@ fn request() -> TurnRequest { request } +#[test] +fn estimates_turn_request_with_system_tools_messages_and_output() { + let mut request = TurnRequest::new( + "test-model", + vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("12345678".to_string()), + }], + ); + request.system_prompt = Some("1234".to_string()); + request.tools = vec![ToolDefinition { + name: "tool".to_string(), + description: "description".to_string(), + input_schema: serde_json::json!({"type": "object"}), + }]; + request.max_output_tokens = Some(10); + + let expected_input = (request.system_prompt.as_deref().unwrap().len() + + serde_json::to_string(&request.messages[0]).unwrap().len() + + serde_json::to_string(&request.tools[0]).unwrap().len()) as u64; + assert_eq!( + estimate_turn_request_tokens(&request), + expected_input.div_ceil(ESTIMATED_CHARS_PER_TOKEN) + 10 + ); +} + +#[tokio::test] +async fn proactively_compacts_before_starting_an_oversized_request() { + let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); + let mut coordinator = coordinator(runtime.clone()); + coordinator.set_max_context_tokens(Some(1)); + let (_sender, control) = turn_control(); + + let block = coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap(); + + assert_eq!(block, ProviderRunBlock::ContextWindowExceeded); + assert!(runtime.requests().is_empty()); +} + fn started(id: &str) -> ScriptedEvent { Ok(AgentEvent::TurnStarted { runtime_request_id: id.to_string(), @@ -1166,6 +1208,7 @@ async fn transcript_projector_emits_one_ui_stream_for_the_whole_run() { max_context_tokens: Some(100_000), capabilities: RuntimeCapabilities::provider(), empty_output_message: None, + todo_items: None, }); let mut ui_events = Vec::new(); let (_sender, control) = turn_control(); @@ -1224,6 +1267,7 @@ fn transcript_projector_preserves_provider_failure_message() { max_context_tokens: Some(100_000), capabilities: RuntimeCapabilities::provider(), empty_output_message: None, + todo_items: None, }); let events = projector .finish( diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index 310d702d..818c2e93 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -21,9 +21,9 @@ use super::rig_tool::action_from_tool_call; use super::ProviderRunProfile; use crate::ai::agent::api::RequestParams; use crate::ai::agent::AIAgentAction; -use crate::ai::bedrock::client::BedrockClient; -use crate::ai::bedrock::convert::CachingConfig; -use crate::ai::bedrock::external_config::ExternalBedrockConfig; +use crate::ai::provider::client::BedrockClient; +use crate::ai::provider::convert::CachingConfig; +use crate::ai::provider::external_config::ExternalBedrockConfig; use crate::ai::provider::types::ConversationMessage; use crate::ai::runtime::RuntimeResponseConfig; use crate::settings::OpenAIProviderKind; @@ -137,6 +137,7 @@ pub(crate) async fn prepare_provider_run( task_id, needs_create_task, user_query, + todo_items, request, persistent_messages, tool_result_archive, @@ -159,6 +160,7 @@ pub(crate) async fn prepare_provider_run( max_context_tokens, capabilities: base_runtime.descriptor().capabilities.clone(), empty_output_message: None, + todo_items, }; Ok(PreparedProviderRun { base_profile: ProviderRunProfile::new(base_runtime, request), diff --git a/app/src/ai/runtime/rig_request.rs b/app/src/ai/runtime/rig_request.rs index 85bf6a7d..1180afc2 100644 --- a/app/src/ai/runtime/rig_request.rs +++ b/app/src/ai/runtime/rig_request.rs @@ -11,21 +11,23 @@ use galaxy_agent_core::{ }; use sha2::{Digest as _, Sha256}; use uuid::Uuid; +use warp_multi_agent_api as api; use warp_multi_agent_api::ToolType; use crate::ai::agent::api::RequestParams; use crate::ai::agent::{AIAgentContext, AIAgentInput, MCPContext, UserQueryMode}; -use crate::ai::bedrock::request_translator::{ - default_tool_definitions, sanitize_messages_for_bedrock, tool_name_is_supported, -}; use crate::ai::openai::client::OpenAIClientConfig; use crate::ai::openai::request_translator::sanitize_messages_for_openai; +use crate::ai::provider::request_translator::{ + default_tool_definitions, sanitize_messages_for_bedrock, tool_name_is_supported, +}; use crate::ai::provider::types::flatten_tool_history_for_no_tools_turn; pub(crate) struct PreparedRigTurn { pub task_id: String, pub needs_create_task: bool, pub user_query: Option, + pub todo_items: Option>, pub request: TurnRequest, pub persistent_messages: Vec, pub tool_result_archive: Vec, @@ -226,6 +228,8 @@ fn prepare_rig_turn_for_provider( .. } = params; + let todo_items = todo_items_from_tasks(&tasks); + let task_id = root_task_id .or_else(|| tasks.first().map(|task| task.id.clone())) .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); @@ -299,6 +303,7 @@ fn prepare_rig_turn_for_provider( task_id, needs_create_task, user_query, + todo_items, request, persistent_messages, tool_result_archive, @@ -307,6 +312,29 @@ fn prepare_rig_turn_for_provider( } } +// Reuse a plan emitted by the model when the existing transcript contains one. This keeps +// direct-provider projection aligned with UpdateTodos instead of inventing a second plan. +fn todo_items_from_tasks(tasks: &[api::Task]) -> Option> { + let mut items = None; + for task in tasks { + for message in &task.messages { + let Some(api::message::Message::UpdateTodos(update)) = &message.message else { + continue; + }; + match update.operation.as_ref()? { + api::message::update_todos::Operation::CreateTodoList(create) => { + items = Some(create.initial_todos.clone()); + } + api::message::update_todos::Operation::UpdatePendingTodos(update) => { + items = Some(update.updated_pending_todos.clone()); + } + api::message::update_todos::Operation::MarkTodosCompleted(_) => {} + } + } + } + items.filter(|items| !items.is_empty()) +} + fn input_messages( inputs: Vec, tool_results: Vec, diff --git a/app/src/ai/runtime/rig_request_tests.rs b/app/src/ai/runtime/rig_request_tests.rs index a72a2cb8..d8e434c4 100644 --- a/app/src/ai/runtime/rig_request_tests.rs +++ b/app/src/ai/runtime/rig_request_tests.rs @@ -32,7 +32,6 @@ fn config() -> OpenAIClientConfig { reasoning_effort: None, max_input_tokens: Some(128_000), max_output_tokens: Some(8_192), - use_rig: true, supports_system_messages: true, } } diff --git a/app/src/ai/runtime/rig_tests.rs b/app/src/ai/runtime/rig_tests.rs index cc2cda25..bf4732ed 100644 --- a/app/src/ai/runtime/rig_tests.rs +++ b/app/src/ai/runtime/rig_tests.rs @@ -16,7 +16,6 @@ fn openai_config(model: &str) -> OpenAIClientConfig { reasoning_effort: None, max_input_tokens: Some(128_000), max_output_tokens: Some(8_192), - use_rig: true, supports_system_messages: true, } } diff --git a/app/src/ai/runtime/rig_tool.rs b/app/src/ai/runtime/rig_tool.rs index ac0e4b9e..94bfcb31 100644 --- a/app/src/ai/runtime/rig_tool.rs +++ b/app/src/ai/runtime/rig_tool.rs @@ -102,7 +102,7 @@ pub(crate) fn action_from_tool_call( optional_bounded_u64( input, "wait_seconds", - crate::ai::bedrock::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS, + crate::ai::provider::request_translator::COMMAND_MONITOR_MAX_POLL_SECONDS, )? .unwrap_or(2), ))), diff --git a/app/src/code_review/code_review_header/mod.rs b/app/src/code_review/code_review_header/mod.rs index 807e6484..3db1443a 100644 --- a/app/src/code_review/code_review_header/mod.rs +++ b/app/src/code_review/code_review_header/mod.rs @@ -297,7 +297,8 @@ impl CodeReviewHeader { .finish(), ); } - right_section_compact.add_child(ChildView::new(&code_review_header_fields.diff_selector).finish()); + right_section_compact + .add_child(ChildView::new(&code_review_header_fields.diff_selector).finish()); right_section_compact.add_child(Container::new(right_subsection_compact.finish()).finish()); Clipped::new( diff --git a/app/src/settings/ai.rs b/app/src/settings/ai.rs index fbb8ce97..7848451e 100644 --- a/app/src/settings/ai.rs +++ b/app/src/settings/ai.rs @@ -828,11 +828,6 @@ pub struct BedrockModelConfig { #[serde(default)] #[schemars(description = "Whether the model supports image/vision input.")] pub vision_supported: bool, - #[serde(default)] - #[schemars( - description = "Route this model through Galaxy's Rig Bedrock runtime. Disabled by default while compatibility validation is in progress." - )] - pub use_rig: bool, } impl settings_value::SettingsValue for BedrockModelConfig {} @@ -883,11 +878,6 @@ pub struct OpenAIModelConfig { description = "Optional provider hint (e.g. anthropic, openai, google) for icon display." )] pub provider: Option, - #[serde(default)] - #[schemars( - description = "Route this model through Galaxy's Rig runtime. This is an opt-in migration path." - )] - pub use_rig: bool, #[serde(default, skip_serializing_if = "Option::is_none")] #[schemars( description = "Whether this endpoint accepts system-role messages. Set false for ChatGPT-backed LiteLLM models that reject them." diff --git a/app/src/settings/ai_tests.rs b/app/src/settings/ai_tests.rs index efafb541..367bcfc2 100644 --- a/app/src/settings/ai_tests.rs +++ b/app/src/settings/ai_tests.rs @@ -444,7 +444,6 @@ fn orchestration_is_enabled_when_ai_is_enabled() { model_id: "test-model".to_string(), display_name: "Test model".to_string(), vision_supported: false, - use_rig: false, }], ctx, ) diff --git a/app/src/settings_view/ai_page.rs b/app/src/settings_view/ai_page.rs index e754f113..2ee3c93e 100644 --- a/app/src/settings_view/ai_page.rs +++ b/app/src/settings_view/ai_page.rs @@ -7251,9 +7251,6 @@ impl ProviderSettingsWidget { if model.effective_vision_supported() { details.push("Images".to_string()); } - if model.use_rig { - details.push("Rig".to_string()); - } if !model.reasoning_efforts.is_empty() { details.push(format!("Reasoning: {}", model.reasoning_efforts.join(", "))); } @@ -7361,9 +7358,6 @@ impl ProviderSettingsWidget { if model.vision_supported { details.push("Images".to_string()); } - if model.use_rig { - details.push("Rig".to_string()); - } Self::render_model_row( model.display_name.clone(), model.model_id.clone(), diff --git a/app/src/settings_view/provider_setup_view.rs b/app/src/settings_view/provider_setup_view.rs index 831c6633..575d2b27 100644 --- a/app/src/settings_view/provider_setup_view.rs +++ b/app/src/settings_view/provider_setup_view.rs @@ -949,7 +949,7 @@ impl ProviderSetupView { return; } ProviderSetupProviderType::Bedrock => { - let config = crate::ai::bedrock::client::BedrockClientConfig { + let config = crate::ai::provider::client::BedrockClientConfig { auth_method: self.draft_bedrock.auth_method, profile: self.draft_bedrock.profile.clone(), region: self.draft_bedrock.region.clone(), @@ -957,11 +957,10 @@ impl ProviderSetupView { secret_access_key: self.draft_bedrock.secret_access_key.clone(), session_token: None, cross_region_inference: self.draft_bedrock.cross_region_inference, - use_rig: false, }; ctx.spawn( async move { - crate::ai::bedrock::discovery::discover_available_models(config).await + crate::ai::provider::discovery::discover_available_models(config).await }, move |me, result, ctx| match result { Ok(models) => { diff --git a/app/src/terminal/input/slash_commands/mod.rs b/app/src/terminal/input/slash_commands/mod.rs index fc027fa5..462f8a06 100644 --- a/app/src/terminal/input/slash_commands/mod.rs +++ b/app/src/terminal/input/slash_commands/mod.rs @@ -1189,21 +1189,21 @@ impl Input { .enumerate() .map(|(i, msg)| { let content_str = match &msg.content { - crate::ai::bedrock::convert::MessageContent::Text(t) => { + crate::ai::provider::convert::MessageContent::Text(t) => { if t.len() > 500 { format!("{}... ({} chars total)", &t[..500], t.len()) } else { t.clone() } } - crate::ai::bedrock::convert::MessageContent::ToolUse { + crate::ai::provider::convert::MessageContent::ToolUse { name, tool_use_id, .. } => { format!("ToolUse(name={name}, id={tool_use_id})") } - crate::ai::bedrock::convert::MessageContent::ToolResult { + crate::ai::provider::convert::MessageContent::ToolResult { tool_use_id, content, is_error, @@ -1217,7 +1217,7 @@ impl Input { "ToolResult(id={tool_use_id}, err={is_error}): {truncated}" ) } - crate::ai::bedrock::convert::MessageContent::MultiPart(parts) => { + crate::ai::provider::convert::MessageContent::MultiPart(parts) => { format!("MultiPart({} parts)", parts.len()) } }; diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index 33535467..91846d73 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -6994,7 +6994,7 @@ impl TerminalView { } fn toggle_settings_view(&mut self, ctx: &mut ViewContext) { - use crate::ai::bedrock::settings_view::SettingsView; + use crate::ai::provider::settings_view::SettingsView; // If already showing, remove it if let Some(view_id) = self.settings_view_id.take() { @@ -10912,7 +10912,7 @@ impl TerminalView { // Use the configured command, but if it's the bare default ("aws sso login") // prefer the external config's auth_refresh_command which includes --profile. let login_command = if settings_command == "aws sso login" { - use crate::ai::bedrock::external_config::ExternalBedrockConfig; + use crate::ai::provider::external_config::ExternalBedrockConfig; let external = ExternalBedrockConfig::load(); external.auth_refresh_command.unwrap_or(settings_command) } else { @@ -10948,34 +10948,37 @@ impl TerminalView { |_me, result, ctx| { match result { Ok(()) => { - log::info!("[bedrock] AWS login completed successfully, refreshing credentials and resuming"); - // Refresh credentials from the updated SSO cache - ApiKeyManager::handle(ctx).update( - ctx, - |manager, ctx| { - drop(crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx)); - }, - ); - // Resume the conversation after a short delay to let credentials load - let _ = ctx.spawn( - async move { - tokio::time::sleep(std::time::Duration::from_millis(500)).await; - }, - |me, _, ctx| { - // Find the active conversation and resume it directly - let conversation_id = if FeatureFlag::AgentView.is_enabled() { - me.agent_view_controller - .as_ref(ctx) - .agent_view_state() - .active_conversation_id() - } else { - BlocklistAIHistoryModel::as_ref(ctx).last_conversation_id(me.id()) - }; - if let Some(conversation_id) = conversation_id { - me.handle_resume_conversation(&conversation_id, ctx); + log::info!("[bedrock] AWS login completed successfully; reloading credentials before resuming"); + + // Refresh through ApiKeyManager and wait for its completion. A fixed + // delay can resume the request while the SDK still has stale SSO data. + let refresh = ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| { + crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx) + }); + let _ = ctx.spawn(refresh, |me, result, ctx| { + match result { + Ok(()) => { + log::info!("[bedrock] AWS credentials refreshed; resuming conversation"); + let conversation_id = if FeatureFlag::AgentView.is_enabled() { + me.agent_view_controller + .as_ref(ctx) + .agent_view_state() + .active_conversation_id() + } else { + BlocklistAIHistoryModel::as_ref(ctx) + .last_conversation_id(me.id()) + }; + if let Some(conversation_id) = conversation_id { + me.handle_resume_conversation(&conversation_id, ctx); + } } - }, - ); + Err(error) => { + log::error!( + "[bedrock] AWS login succeeded but credential reload failed; not resuming: {error}" + ); + } + } + }); } Err(e) => { log::error!("[bedrock] AWS login failed: {e}"); @@ -20646,6 +20649,35 @@ impl TerminalView { AIBlockEvent::RunAwsLoginCommand => { self.run_aws_login_command(ctx); } + AIBlockEvent::RefreshAwsCredentials => { + let refresh = ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| { + crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx) + }); + let _ = ctx.spawn(refresh, |me, result, ctx| { + match result { + Ok(()) => { + log::info!("[bedrock] AWS credentials reloaded; resuming conversation"); + let conversation_id = if FeatureFlag::AgentView.is_enabled() { + me.agent_view_controller + .as_ref(ctx) + .agent_view_state() + .active_conversation_id() + } else { + BlocklistAIHistoryModel::as_ref(ctx).last_conversation_id(me.id()) + }; + if let Some(conversation_id) = conversation_id { + me.handle_resume_conversation(&conversation_id, ctx); + } + } + Err(error) => { + log::warn!( + "[bedrock] AWS credential reload failed; starting configured login command: {error}" + ); + me.run_aws_login_command(ctx); + } + } + }); + } } ctx.notify(); } diff --git a/crates/galaxy_agent_core/src/provider_run.rs b/crates/galaxy_agent_core/src/provider_run.rs index 64b0f317..b76fc4b9 100644 --- a/crates/galaxy_agent_core/src/provider_run.rs +++ b/crates/galaxy_agent_core/src/provider_run.rs @@ -1010,6 +1010,29 @@ impl ProviderRun { Ok(ModelFailureDisposition::RunFailed) } + /// Parks a failed model call so its transcript can be compacted before retrying it. + /// The run and current work identity are intentionally unchanged. + pub fn prepare_context_compaction( + &mut self, + work_id: &ExternalWorkId, + ) -> Result<(), ProviderRunProtocolError> { + let call = match &self.state { + ProviderRunState::AwaitingModel { call } => call, + ProviderRunState::ReadyToCallModel + | ProviderRunState::ResolvingModel { .. } + | ProviderRunState::AwaitingTools { .. } + | ProviderRunState::AwaitingDriver { .. } + | ProviderRunState::Done { .. } + | ProviderRunState::Failed { .. } + | ProviderRunState::Cancelled { .. } => { + return Err(self.unexpected_state(ProviderRunPhase::AwaitingModel)); + } + }; + validate_work_id(&call.work_id, work_id)?; + self.state = ProviderRunState::ReadyToCallModel; + Ok(()) + } + pub fn request_tool_permission( &mut self, work_id: &ExternalWorkId, diff --git a/crates/galaxy_agent_core/src/provider_run_tests.rs b/crates/galaxy_agent_core/src/provider_run_tests.rs index dd43336e..854894da 100644 --- a/crates/galaxy_agent_core/src/provider_run_tests.rs +++ b/crates/galaxy_agent_core/src/provider_run_tests.rs @@ -1265,6 +1265,19 @@ fn transcript_compaction_is_rejected_during_a_model_call() { )); } +#[test] +fn context_compaction_preserves_run_and_work_identity() { + let mut run = run(); + let call = next_model_call(&mut run); + let run_id = run.id().clone(); + + run.prepare_context_compaction(&call.work_id).unwrap(); + + assert_eq!(run.id(), &run_id); + assert_eq!(run.ready_work_id(), Some(call.work_id)); + assert_eq!(run.state().phase(), ProviderRunPhase::ReadyToCallModel); +} + #[test] fn restore_normalization_commits_a_fully_resolved_batch() { let mut run = run(); diff --git a/crates/galaxy_agent_rig/src/stream.rs b/crates/galaxy_agent_rig/src/stream.rs index 65328753..d28236cd 100644 --- a/crates/galaxy_agent_rig/src/stream.rs +++ b/crates/galaxy_agent_rig/src/stream.rs @@ -392,10 +392,12 @@ fn completion_error_indicates_recoverable_transport(error: &CompletionError) -> // Rig flattens non-status SSE transport failures into ProviderError. // Preserve their transport semantics so the durable provider run can // retry a truncated or otherwise interrupted response stream. - CompletionError::ProviderError(message) => message - .to_ascii_lowercase() - .starts_with("http client error:"), - _ => false, + CompletionError::ProviderError(message) => { + let normalized = message.to_ascii_lowercase(); + normalized.starts_with("http client error:") + || text_indicates_interrupted_stream(&normalized) + } + _ => text_indicates_interrupted_stream(&error.to_string().to_ascii_lowercase()), } } @@ -414,10 +416,29 @@ fn text_indicates_transient_provider_failure(text: &str) -> bool { || normalized.contains("service unavailable") || normalized.contains("server is overloaded") || normalized.contains("servers are currently overloaded") + || normalized.contains("overload") + || normalized.contains("bad gateway") + || normalized.contains("gateway timeout") + || normalized.contains("upstream connect error") + || normalized.contains("upstream request timeout") + || normalized.contains("upstream timed out") + || normalized.contains("upstream unavailable") +} + +fn text_indicates_interrupted_stream(text: &str) -> bool { + text.contains("unexpected eof") + || text.contains("unexpected end of file") + || text.contains("stream terminated") + || text.contains("stream ended") + || text.contains("stream closed") + || text.contains("connection reset") + || text.contains("connection closed") + || text.contains("incomplete message") } fn map_completion_error(error: CompletionError) -> AgentError { let is_context_window_exceeded = completion_error_indicates_context_window_exceeded(&error); + let is_authentication_failure = completion_error_indicates_authentication_failure(&error); let is_recoverable_transport = completion_error_indicates_recoverable_transport(&error); let is_transient_provider_failure = completion_error_indicates_transient_provider_failure(&error); @@ -426,7 +447,13 @@ fn map_completion_error(error: CompletionError) -> AgentError { .map(|status| status.as_u16()); let kind = if is_context_window_exceeded { AgentErrorKind::ContextWindowExceeded - } else if is_recoverable_transport { + } else if is_authentication_failure { + // Bedrock sometimes returns credential errors as a provider error with no + // HTTP status (and Rig can flatten the response body). Classify by the + // service error code as well as status so the host can run its existing + // credential refresh/login flow. + AgentErrorKind::Authentication + } else if is_recoverable_transport || matches!(status, Some(408 | 425)) { AgentErrorKind::Transport } else { match status { @@ -450,14 +477,47 @@ fn map_completion_error(error: CompletionError) -> AgentError { } }; let mut mapped = AgentError::new(kind, error.to_string()); + if is_authentication_failure { + mapped.user_message = Some( + "AWS credentials for Bedrock are missing or expired. Refresh your AWS login and try again." + .to_string(), + ); + } mapped.recoverable = matches!( kind, AgentErrorKind::RateLimited | AgentErrorKind::Transport - ) || status.is_some_and(|status| (500..=599).contains(&status)) - || is_transient_provider_failure; + ) || status.is_some_and(|status| { + matches!(status, 408 | 425 | 502 | 503 | 504) || (500..=599).contains(&status) + }) || is_transient_provider_failure; mapped } +fn completion_error_indicates_authentication_failure(error: &CompletionError) -> bool { + error + .provider_response_body() + .is_some_and(text_indicates_authentication_failure) + || text_indicates_authentication_failure(&error.to_string()) +} + +fn text_indicates_authentication_failure(text: &str) -> bool { + let normalized = text.to_ascii_lowercase(); + [ + "accessdenied", + "access denied", + "expiredtoken", + "expired token", + "invalidclienttokenid", + "invalid client token", + "unrecognizedclient", + "unrecognized client", + "security token included in the request is invalid", + "unable to locate credentials", + "credential should be scoped to a valid region", + ] + .iter() + .any(|marker| normalized.contains(marker)) +} + #[cfg(test)] mod tests { use galaxy_agent_core::{AgentErrorKind, StopReason}; @@ -560,6 +620,20 @@ mod tests { assert_eq!(mapped.kind, AgentErrorKind::Provider); assert!(mapped.recoverable); } + + #[test] + fn flattened_bedrock_credential_errors_are_authentication_failures() { + for message in [ + r#"{"__type":"ExpiredTokenException","message":"The security token included in the request is expired"}"#, + "UnrecognizedClientException: The security token included in the request is invalid", + "AccessDeniedException: not authorized to perform bedrock:ConverseStream", + ] { + let mapped = map_completion_error(CompletionError::ProviderError(message.to_string())); + assert_eq!(mapped.kind, AgentErrorKind::Authentication); + assert!(!mapped.recoverable); + assert!(mapped.user_message.is_some()); + } + } } #[cfg(test)] diff --git a/crates/galaxy_agent_rig/src/stream_tests.rs b/crates/galaxy_agent_rig/src/stream_tests.rs index df823873..61df2a33 100644 --- a/crates/galaxy_agent_rig/src/stream_tests.rs +++ b/crates/galaxy_agent_rig/src/stream_tests.rs @@ -27,3 +27,58 @@ fn flattened_streamed_provider_overload_is_recoverable() { assert_eq!(mapped.kind, AgentErrorKind::Provider); assert!(mapped.recoverable); } + +#[test] +fn transient_http_statuses_are_recoverable() { + for status_code in [408, 425, 502, 503, 504] { + let status = rig_core::http_client::Response::builder() + .status(status_code) + .body(()) + .unwrap() + .status(); + let error = CompletionError::from_http_response(status, "transient failure"); + + let mapped = map_completion_error(error); + + assert!(mapped.recoverable, "HTTP {status_code} should be retryable"); + assert_eq!( + mapped.kind, + if matches!(status_code, 408 | 425) { + AgentErrorKind::Transport + } else { + AgentErrorKind::Provider + }, + "unexpected classification for HTTP {status_code}" + ); + } +} + +#[test] +fn transient_gateway_upstream_and_overload_messages_are_recoverable() { + for message in [ + "502 Bad Gateway", + "upstream request timeout", + "upstream timed out", + "The server is overloaded; retry later", + ] { + let mapped = map_completion_error(CompletionError::ProviderError(message.to_string())); + + assert_eq!(mapped.kind, AgentErrorKind::Provider); + assert!(mapped.recoverable, "message should be retryable: {message}"); + } +} + +#[test] +fn interrupted_stream_messages_are_recoverable_transport() { + for message in [ + "Http client error: unexpected EOF", + "Http client error: stream terminated unexpectedly", + "Http client error: connection reset by peer", + "unexpected EOF", + ] { + let mapped = map_completion_error(CompletionError::ProviderError(message.to_string())); + + assert_eq!(mapped.kind, AgentErrorKind::Transport); + assert!(mapped.recoverable, "message should be retryable: {message}"); + } +} diff --git a/crates/integration/src/test/rig_runtime.rs b/crates/integration/src/test/rig_runtime.rs index ee8f6b7a..cad57e22 100644 --- a/crates/integration/src/test/rig_runtime.rs +++ b/crates/integration/src/test/rig_runtime.rs @@ -383,7 +383,6 @@ base_url = "http://{address}/v1" model_id = "{MODEL_ID}" display_name = "Rig Integration Model" context_size = 128000 -use_rig = true supports_system_messages = false "# ); diff --git a/migration-docs/project-architecture.md b/migration-docs/project-architecture.md index 0237d54d..c1eb73cb 100644 --- a/migration-docs/project-architecture.md +++ b/migration-docs/project-architecture.md @@ -18,7 +18,7 @@ Galaxy is a fork of Warp (the terminal emulator) rebranded under Samsung/Ryan Wa ## Critical Architecture Decisions -1. **AI flows through Bedrock exclusively** — `app/src/ai/bedrock/client.rs` (`BedrockClient::converse_stream`) +1. **AI flows through Bedrock exclusively** — `app/src/ai/provider/client.rs` (`BedrockClient::converse_stream`) 2. **Models:** Default = Claude Opus 4.6 on Bedrock; also supports Claude Sonnet, Haiku, Nova Pro/Lite/Micro, DeepSeek R1 3. **Bedrock auth methods:** AWS Profile (default), Static Keys, SSO 4. **Cross-region inference:** Auto-prefixes model IDs (us., eu., jp., apac., au.) based on configured region @@ -31,7 +31,7 @@ Galaxy is a fork of Warp (the terminal emulator) rebranded under Samsung/Ryan Wa - `app/` — Main binary and application logic - `app/src/ai/` — AI module (agent, bedrock, mcp, skills, ambient agents, etc.) - - `app/src/ai/bedrock/` — Direct Bedrock client (client.rs, convert.rs, stream.rs, models.rs, convert_request.rs) + - `app/src/ai/provider/` — Direct Bedrock client (client.rs, convert.rs, stream.rs, models.rs, convert_request.rs) - `app/src/settings/ai.rs` — AI settings (Bedrock config, permissions, autoexecution rules) - `app/src/terminal/` — Terminal emulation - `app/src/code/` — Code editor mode diff --git a/plans/galaxy-refactor.md b/plans/galaxy-refactor.md index 1ac6f0b6..d55d1b4e 100644 --- a/plans/galaxy-refactor.md +++ b/plans/galaxy-refactor.md @@ -273,10 +273,10 @@ These are bugs in the current branch that must be addressed regardless of phase: | 🟠 HIGH | Windows registry login item `"Warp"` keys | `app/src/login_item/windows.rs` | 3c | | 🟠 HIGH | Config migration missing `.warp` origin | `crates/warp_core/src/paths.rs` | 6e | | 🟠 HIGH | `needs_create_task` false negative | `app/src/ai/agent/api/impl.rs` | 6f | -| 🟡 MEDIUM | `TEXT_FLUSH_THRESHOLD` text discard | `app/src/ai/bedrock/stream.rs` | 6b | -| 🟡 MEDIUM | Empty synthesized tool results | `app/src/ai/bedrock/convert_request.rs` | 6c | -| 🟡 MEDIUM | 9-param `converse_stream()` | `app/src/ai/bedrock/client.rs` | 6a | -| 🟡 MEDIUM | Diagnostic logger privacy | `app/src/ai/bedrock/diagnostic.rs` | 6d | +| 🟡 MEDIUM | `TEXT_FLUSH_THRESHOLD` text discard | `app/src/ai/provider/stream.rs` | 6b | +| 🟡 MEDIUM | Empty synthesized tool results | `app/src/ai/provider/convert_request.rs` | 6c | +| 🟡 MEDIUM | 9-param `converse_stream()` | `app/src/ai/provider/client.rs` | 6a | +| 🟡 MEDIUM | Diagnostic logger privacy | `app/src/ai/provider/diagnostic.rs` | 6d | ---