[PR #247] [MERGED] Runtime provider switching #267

Closed
opened 2026-06-06 22:09:57 -04:00 by yindo · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/vxcontrol/pentagi/pull/247
Author: @asdek
Created: 4/8/2026
Status: Merged
Merged: 4/8/2026
Merged by: @asdek

Base: mainHead: feature/switch_provider


📝 Commits (1)

  • 7c25f35 feat: enable runtime provider switching

📊 Changes

52 files changed (+594 additions, -196 deletions)

View changed files

📝 backend/cmd/ctester/main.go (+10 -10)
📝 backend/pkg/controller/context.go (+3 -2)
📝 backend/pkg/controller/flow.go (+65 -3)
📝 backend/pkg/controller/subtask.go (+4 -0)
📝 backend/pkg/database/flows.sql.go (+43 -0)
📝 backend/pkg/database/msgchains.sql.go (+11 -3)
📝 backend/pkg/database/querier.go (+1 -0)
📝 backend/pkg/graph/generated.go (+31 -4)
📝 backend/pkg/graph/schema.graphqls (+1 -1)
📝 backend/pkg/graph/schema.resolvers.go (+24 -4)
📝 backend/pkg/providers/anthropic/anthropic.go (+11 -1)
📝 backend/pkg/providers/anthropic/anthropic_test.go (+2 -2)
📝 backend/pkg/providers/assistant.go (+17 -4)
📝 backend/pkg/providers/bedrock/bedrock.go (+11 -1)
📝 backend/pkg/providers/bedrock/bedrock_test.go (+14 -14)
📝 backend/pkg/providers/custom/custom.go (+11 -1)
📝 backend/pkg/providers/custom/custom_test.go (+3 -3)
📝 backend/pkg/providers/custom/example_test.go (+3 -2)
📝 backend/pkg/providers/deepseek/deepseek.go (+11 -1)
📝 backend/pkg/providers/deepseek/deepseek_test.go (+7 -7)

...and 32 more files

📄 Description

Description of the Change

Problem

Two related but distinct problems affecting flow execution reliability and flexibility:

Message Chain Corruption on Server Restart:

  • When the backend was restarted while a subtask was actively running, the persisted message chain in the database contained AI messages with pending tool calls that had no corresponding tool responses
  • On restart, LoadSubtaskWorker would resume execution using this corrupted chain directly via PerformAgentChain without any consistency check
  • This caused API errors from all providers: Kimi/Moonshot reported "tool_call_ids did not have response messages", Anthropic reported "tool_use ids were found without tool_result blocks"
  • The reflector agent failed to recover because the error was at the chain level, not the response level
  • Users had no way to resume a broken flow without manual intervention

No Runtime Provider Switching:

  • Once a flow was created with a specific LLM provider, it was permanently bound to that provider for all tasks and subtasks
  • Users could not change the provider mid-flight if the current provider was unavailable, rate-limited, or simply not optimal for a particular phase of the penetration test
  • When providers were switched, the existing message chains contained provider-specific tool call ID formats (e.g., Kimi uses {f}:{r:1:d} format like coder:14, Anthropic uses toolu_{r:24:b}) and reasoning content signatures that would be rejected by a different provider
  • No API existed to request a provider switch alongside user input

Solution

Conditional Chain Normalization in processChain:

The core fix is a single, well-placed normalization step in processChain — the shared method that handles all message chain database operations. Before saving an updated chain back to the database, the method now compares msgChain.ModelProvider and msgChain.Model (stored in the database record) against the current flowProvider's active provider and model.

  • If provider or model has not changed: the chain is saved as-is, preserving all reasoning content and signatures for LLM-side KV cache efficiency
  • If provider or model has changed: the chain undergoes full normalization — NormalizeToolCallIDs converts all tool call IDs to the new provider's format, ClearReasoning removes provider-specific signatures, and if the last message is not a Human message, a continuation message ("Continue from where we left off") is appended to ensure the new provider receives a clean conversation boundary

This approach means chain normalization happens exactly once — the next time a user adds input and processChain is called — rather than unconditionally on every PerformAgentChain call, which would destroy reasoning cache.

Pre-execution Chain Consistency:

EnsureChainConsistency is called at the start of subtaskWorker.Run, before PerformAgentChain. This uses NewChainAST(force=true) to add fallback tool responses for any pending tool calls left over from an interrupted run. Combined with the normalization in processChain, this covers the restart scenario: the chain is fixed when the user resumes, not when the backend blindly retries.

Runtime Provider Switch via GraphQL:

The putUserInput mutation now accepts an optional modelProvider parameter. When provided, switchProvider is called before the input is enqueued. switchProvider calls the existing SetProvider method on the current flowProvider — which atomically updates both the underlying provider instance and the tcIDTemplate under a mutex — and then persists the change to the flows table via a new UpdateFlowProvider SQL query. The executor and containers are reused entirely; no restart or reinitialization is needed.

UpdateMsgChain with Provider Metadata:

The UpdateMsgChain SQL query is extended to update model and model_provider alongside the chain content. This ensures the stored metadata always reflects the provider that last wrote the chain, enabling accurate change detection on the next processChain call.

Frontend Provider Selection:

The automation message form now allows provider selection when the flow is in waiting status. The provider button is enabled only in waiting state and disabled in all other states (running, created, finished, failed). The selected provider name is passed to the putUserInput mutation.

Closes #

Type of Change

  • 🐛 Bug fix (non-breaking change which fixes an issue)
  • 🚀 New feature (non-breaking change which adds functionality)
  • 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • 📚 Documentation update
  • 🔧 Configuration change
  • 🧪 Test update
  • 🛡️ Security update

Areas Affected

  • Core Services (Frontend UI/Backend API)
  • AI Agents (Researcher/Developer/Executor)
  • Security Tools Integration
  • Memory System (Vector Store/Knowledge Base)
  • Monitoring Stack (Grafana/OpenTelemetry)
  • Analytics Platform (Langfuse)
  • External Integrations (LLM/Search APIs)
  • Documentation
  • Infrastructure/DevOps

Testing and Verification

Test Configuration

PentAGI Version: Latest development
Docker Version: 24.0.x+
Host OS: macOS/Linux
LLM Provider: Kimi (moonshot/kimi-k2.5), Anthropic (claude-sonnet), OpenAI
Enabled Features: Core

Test Steps

  1. Restart recovery: Start a flow with Kimi, let it reach an active subtask, restart the backend container mid-execution, verify the flow enters waiting state with an error rather than looping on API 400 errors
  2. Resume after restart: After restart recovery above, type resume in the message input, verify the subtask continues without API errors and completes successfully
  3. Provider switch mid-flow: Start a flow with Kimi, wait for it to enter waiting state, select a different provider (e.g. Anthropic) in the provider dropdown, submit a message — verify the flow continues with the new provider and no tool call ID format errors occur
  4. Provider dropdown state: Verify the provider button is disabled when flow is running or created, and enabled when flow is waiting
  5. Cache preservation: Start a flow without switching providers, submit a user response to ask_user, verify no unnecessary chain rewriting in logs (absence of "provider or model changed" log line)
  6. Reasoning cleared on switch: Switch from a reasoning-capable provider (Kimi) to one without reasoning (OpenAI), verify no reasoning signatures appear in the chain sent to OpenAI

Test Results

  • Restart recovery confirmed: Flow 331 (Anthropic) and Flow 332 (Kimi) both recovered after backend restart when user submitted resume — previously both looped with 400 errors until manual DB intervention
  • Provider switch functional: Kimi → Anthropic switch mid-flow completed without tool call ID format errors
  • Tool ID normalization: advice:2 (Kimi format) correctly converted to toolu_... (Anthropic format) on provider switch
  • Reasoning cleared: No Kimi reasoning signatures present in chains after switching to OpenAI
  • Cache unaffected: No unnecessary chain rewrites when provider does not change
  • UI state correct: Provider button locked during running, unlocked during waiting
  • Backend compilation clean
  • Frontend build passed

Security Considerations

No Security Impact:

  • Provider switching is scoped to the authenticated user's flow — existing GraphQL permission checks apply to putUserInput
  • The SetProvider call operates entirely within the in-process flowProvider struct, no new network surfaces introduced
  • UpdateFlowProvider updates only model_provider_name, model_provider_type, tool_call_id_template, and model for the authenticated user's flow
  • No new external dependencies

Performance Impact

Improvements:

  • Chain normalization happens only when the provider actually changes — zero overhead on the common path where provider is unchanged
  • LLM KV cache utilization preserved: reasoning content and signatures are retained when the provider does not change, allowing the backend to benefit from cached prefixes
  • EnsureChainConsistency before PerformAgentChain prevents wasted API retries on a chain that would have been rejected anyway

Overhead:

  • One additional GetMsgChain comparison per processChain call (already fetched, no extra DB query)
  • NormalizeToolCallIDs + ClearReasoning + AppendHumanMessage run only on provider change — O(n) over chain messages, typically <1ms

Documentation Updates

  • README.md updates
  • API documentation updates
  • Configuration documentation updates
  • GraphQL schema updates — putUserInput gains optional modelProvider: String parameter
  • Other

Deployment Notes

Database Migration:

  • UpdateMsgChain SQL extended with model and model_provider parameters (sqlc-generated, backward compatible — existing rows retain their current values)
  • New UpdateFlowProvider SQL query added to flows table (no schema change)

Environment Variables:

  • No new configuration required
  • Existing deployments work without changes

Deployment Steps:

  1. Pull latest changes
  2. Rebuild: docker compose build
  3. Restart: docker compose up -d
  4. sqlc code is generated during build

Compatibility:

  • Fully backward compatible — existing flows continue without any changes
  • Flows created before this change will normalize their chains automatically the first time the user submits input after a provider switch
  • No breaking changes to the API — modelProvider is optional; all existing clients continue to work

Checklist

Code Quality

  • My code follows the project's coding standards
  • I have added/updated necessary documentation
  • I have added tests to cover my changes
  • All new and existing tests pass
  • I have run go fmt and go vet (for Go code)
  • I have run npm run lint (for TypeScript/JavaScript code)

Security

  • I have considered security implications
  • Changes maintain or improve the security model
  • Sensitive information has been properly handled

Compatibility

  • Changes are backward compatible
  • Breaking changes are clearly marked and documented
  • Dependencies are properly updated

Documentation

  • Documentation is clear and complete
  • Comments are added for non-obvious code
  • API changes are documented

Additional Notes

Key Design Decisions

Why normalize in processChain rather than PerformAgentChain:
PerformAgentChain runs on every subtask iteration, including iterations where the provider has not changed. Normalizing there would clear reasoning signatures unconditionally, eliminating any KV cache benefit from the LLM backend. processChain is only called when new user input is being stored — exactly the moment a new conversation boundary exists and normalization is safe.

Why add a continuation Human message:
When a provider is switched and the last message in the chain is not a Human message (e.g., the previous turn ended with a tool call whose response was just added), the new provider would receive a chain ending with a Tool message. Some providers reject this or behave unpredictably. Appending "Continue from where we left off" creates a clean conversation turn that any provider can handle correctly without changing the semantic content of the task.

Why reuse the executor on provider switch:
The Docker containers holding pentest tool environments are expensive to create and are unrelated to the LLM provider. SetProvider replaces only the LLM client and tcIDTemplate inside the existing flowProvider struct under a mutex, leaving the executor, image reference, and all log workers intact.

Why the restart scenario still requires user input:
This is intentional. The backend cannot safely resume execution automatically after an unknown interruption — the user should confirm the intent to continue. Once the user submits any input, processChain fires, fixes the chain, and execution resumes. This matches the existing graceful degradation model.


🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/vxcontrol/pentagi/pull/247 **Author:** [@asdek](https://github.com/asdek) **Created:** 4/8/2026 **Status:** ✅ Merged **Merged:** 4/8/2026 **Merged by:** [@asdek](https://github.com/asdek) **Base:** `main` ← **Head:** `feature/switch_provider` --- ### 📝 Commits (1) - [`7c25f35`](https://github.com/vxcontrol/pentagi/commit/7c25f356eda2f3cc18fd921667d6406eee777e2f) feat: enable runtime provider switching ### 📊 Changes **52 files changed** (+594 additions, -196 deletions) <details> <summary>View changed files</summary> 📝 `backend/cmd/ctester/main.go` (+10 -10) 📝 `backend/pkg/controller/context.go` (+3 -2) 📝 `backend/pkg/controller/flow.go` (+65 -3) 📝 `backend/pkg/controller/subtask.go` (+4 -0) 📝 `backend/pkg/database/flows.sql.go` (+43 -0) 📝 `backend/pkg/database/msgchains.sql.go` (+11 -3) 📝 `backend/pkg/database/querier.go` (+1 -0) 📝 `backend/pkg/graph/generated.go` (+31 -4) 📝 `backend/pkg/graph/schema.graphqls` (+1 -1) 📝 `backend/pkg/graph/schema.resolvers.go` (+24 -4) 📝 `backend/pkg/providers/anthropic/anthropic.go` (+11 -1) 📝 `backend/pkg/providers/anthropic/anthropic_test.go` (+2 -2) 📝 `backend/pkg/providers/assistant.go` (+17 -4) 📝 `backend/pkg/providers/bedrock/bedrock.go` (+11 -1) 📝 `backend/pkg/providers/bedrock/bedrock_test.go` (+14 -14) 📝 `backend/pkg/providers/custom/custom.go` (+11 -1) 📝 `backend/pkg/providers/custom/custom_test.go` (+3 -3) 📝 `backend/pkg/providers/custom/example_test.go` (+3 -2) 📝 `backend/pkg/providers/deepseek/deepseek.go` (+11 -1) 📝 `backend/pkg/providers/deepseek/deepseek_test.go` (+7 -7) _...and 32 more files_ </details> ### 📄 Description <!-- Thank you for your contribution to PentAGI! Please fill out this template completely to help us review your changes effectively. Any PR that does not include enough information may be closed at maintainers' discretion. --> ### Description of the Change #### Problem Two related but distinct problems affecting flow execution reliability and flexibility: **Message Chain Corruption on Server Restart:** - When the backend was restarted while a subtask was actively running, the persisted message chain in the database contained AI messages with pending tool calls that had no corresponding tool responses - On restart, `LoadSubtaskWorker` would resume execution using this corrupted chain directly via `PerformAgentChain` without any consistency check - This caused API errors from all providers: Kimi/Moonshot reported `"tool_call_ids did not have response messages"`, Anthropic reported `"tool_use ids were found without tool_result blocks"` - The reflector agent failed to recover because the error was at the chain level, not the response level - Users had no way to resume a broken flow without manual intervention **No Runtime Provider Switching:** - Once a flow was created with a specific LLM provider, it was permanently bound to that provider for all tasks and subtasks - Users could not change the provider mid-flight if the current provider was unavailable, rate-limited, or simply not optimal for a particular phase of the penetration test - When providers were switched, the existing message chains contained provider-specific tool call ID formats (e.g., Kimi uses `{f}:{r:1:d}` format like `coder:14`, Anthropic uses `toolu_{r:24:b}`) and reasoning content signatures that would be rejected by a different provider - No API existed to request a provider switch alongside user input #### Solution **Conditional Chain Normalization in `processChain`:** The core fix is a single, well-placed normalization step in `processChain` — the shared method that handles all message chain database operations. Before saving an updated chain back to the database, the method now compares `msgChain.ModelProvider` and `msgChain.Model` (stored in the database record) against the current `flowProvider`'s active provider and model. - **If provider or model has not changed**: the chain is saved as-is, preserving all reasoning content and signatures for LLM-side KV cache efficiency - **If provider or model has changed**: the chain undergoes full normalization — `NormalizeToolCallIDs` converts all tool call IDs to the new provider's format, `ClearReasoning` removes provider-specific signatures, and if the last message is not a Human message, a continuation message (`"Continue from where we left off"`) is appended to ensure the new provider receives a clean conversation boundary This approach means chain normalization happens exactly once — the next time a user adds input and `processChain` is called — rather than unconditionally on every `PerformAgentChain` call, which would destroy reasoning cache. **Pre-execution Chain Consistency:** `EnsureChainConsistency` is called at the start of `subtaskWorker.Run`, before `PerformAgentChain`. This uses `NewChainAST(force=true)` to add fallback tool responses for any pending tool calls left over from an interrupted run. Combined with the normalization in `processChain`, this covers the restart scenario: the chain is fixed when the user resumes, not when the backend blindly retries. **Runtime Provider Switch via GraphQL:** The `putUserInput` mutation now accepts an optional `modelProvider` parameter. When provided, `switchProvider` is called before the input is enqueued. `switchProvider` calls the existing `SetProvider` method on the current `flowProvider` — which atomically updates both the underlying provider instance and the `tcIDTemplate` under a mutex — and then persists the change to the `flows` table via a new `UpdateFlowProvider` SQL query. The executor and containers are reused entirely; no restart or reinitialization is needed. **UpdateMsgChain with Provider Metadata:** The `UpdateMsgChain` SQL query is extended to update `model` and `model_provider` alongside the chain content. This ensures the stored metadata always reflects the provider that last wrote the chain, enabling accurate change detection on the next `processChain` call. **Frontend Provider Selection:** The automation message form now allows provider selection when the flow is in `waiting` status. The provider button is enabled only in `waiting` state and disabled in all other states (`running`, `created`, `finished`, `failed`). The selected provider name is passed to the `putUserInput` mutation. Closes # ### Type of Change - [x] 🐛 Bug fix (non-breaking change which fixes an issue) - [x] 🚀 New feature (non-breaking change which adds functionality) - [ ] 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] 📚 Documentation update - [ ] 🔧 Configuration change - [ ] 🧪 Test update - [ ] 🛡️ Security update ### Areas Affected - [x] Core Services (Frontend UI/Backend API) - [x] AI Agents (Researcher/Developer/Executor) - [ ] Security Tools Integration - [ ] Memory System (Vector Store/Knowledge Base) - [ ] Monitoring Stack (Grafana/OpenTelemetry) - [ ] Analytics Platform (Langfuse) - [x] External Integrations (LLM/Search APIs) - [ ] Documentation - [ ] Infrastructure/DevOps ### Testing and Verification #### Test Configuration ```yaml PentAGI Version: Latest development Docker Version: 24.0.x+ Host OS: macOS/Linux LLM Provider: Kimi (moonshot/kimi-k2.5), Anthropic (claude-sonnet), OpenAI Enabled Features: Core ``` #### Test Steps 1. **Restart recovery**: Start a flow with Kimi, let it reach an active subtask, restart the backend container mid-execution, verify the flow enters `waiting` state with an error rather than looping on API 400 errors 2. **Resume after restart**: After restart recovery above, type `resume` in the message input, verify the subtask continues without API errors and completes successfully 3. **Provider switch mid-flow**: Start a flow with Kimi, wait for it to enter `waiting` state, select a different provider (e.g. Anthropic) in the provider dropdown, submit a message — verify the flow continues with the new provider and no tool call ID format errors occur 4. **Provider dropdown state**: Verify the provider button is disabled when flow is `running` or `created`, and enabled when flow is `waiting` 5. **Cache preservation**: Start a flow without switching providers, submit a user response to `ask_user`, verify no unnecessary chain rewriting in logs (absence of `"provider or model changed"` log line) 6. **Reasoning cleared on switch**: Switch from a reasoning-capable provider (Kimi) to one without reasoning (OpenAI), verify no reasoning signatures appear in the chain sent to OpenAI #### Test Results - ✅ **Restart recovery confirmed**: Flow 331 (Anthropic) and Flow 332 (Kimi) both recovered after backend restart when user submitted `resume` — previously both looped with 400 errors until manual DB intervention - ✅ **Provider switch functional**: Kimi → Anthropic switch mid-flow completed without tool call ID format errors - ✅ **Tool ID normalization**: `advice:2` (Kimi format) correctly converted to `toolu_...` (Anthropic format) on provider switch - ✅ **Reasoning cleared**: No Kimi reasoning signatures present in chains after switching to OpenAI - ✅ **Cache unaffected**: No unnecessary chain rewrites when provider does not change - ✅ **UI state correct**: Provider button locked during `running`, unlocked during `waiting` - ✅ Backend compilation clean - ✅ Frontend build passed ### Security Considerations **No Security Impact:** - Provider switching is scoped to the authenticated user's flow — existing GraphQL permission checks apply to `putUserInput` - The `SetProvider` call operates entirely within the in-process `flowProvider` struct, no new network surfaces introduced - `UpdateFlowProvider` updates only `model_provider_name`, `model_provider_type`, `tool_call_id_template`, and `model` for the authenticated user's flow - No new external dependencies ### Performance Impact **Improvements:** - Chain normalization happens only when the provider actually changes — zero overhead on the common path where provider is unchanged - LLM KV cache utilization preserved: reasoning content and signatures are retained when the provider does not change, allowing the backend to benefit from cached prefixes - `EnsureChainConsistency` before `PerformAgentChain` prevents wasted API retries on a chain that would have been rejected anyway **Overhead:** - One additional `GetMsgChain` comparison per `processChain` call (already fetched, no extra DB query) - `NormalizeToolCallIDs` + `ClearReasoning` + `AppendHumanMessage` run only on provider change — O(n) over chain messages, typically <1ms ### Documentation Updates - [ ] README.md updates - [ ] API documentation updates - [ ] Configuration documentation updates - [x] GraphQL schema updates — `putUserInput` gains optional `modelProvider: String` parameter - [ ] Other ### Deployment Notes **Database Migration:** - `UpdateMsgChain` SQL extended with `model` and `model_provider` parameters (sqlc-generated, backward compatible — existing rows retain their current values) - New `UpdateFlowProvider` SQL query added to `flows` table (no schema change) **Environment Variables:** - No new configuration required - Existing deployments work without changes **Deployment Steps:** 1. Pull latest changes 2. Rebuild: `docker compose build` 3. Restart: `docker compose up -d` 4. sqlc code is generated during build **Compatibility:** - ✅ Fully backward compatible — existing flows continue without any changes - ✅ Flows created before this change will normalize their chains automatically the first time the user submits input after a provider switch - ✅ No breaking changes to the API — `modelProvider` is optional; all existing clients continue to work ### Checklist #### Code Quality - [x] My code follows the project's coding standards - [ ] I have added/updated necessary documentation - [ ] I have added tests to cover my changes - [x] All new and existing tests pass - [x] I have run `go fmt` and `go vet` (for Go code) - [x] I have run `npm run lint` (for TypeScript/JavaScript code) #### Security - [x] I have considered security implications - [x] Changes maintain or improve the security model - [x] Sensitive information has been properly handled #### Compatibility - [x] Changes are backward compatible - [x] Breaking changes are clearly marked and documented - [x] Dependencies are properly updated #### Documentation - [x] Documentation is clear and complete - [x] Comments are added for non-obvious code - [x] API changes are documented ### Additional Notes ## Key Design Decisions **Why normalize in `processChain` rather than `PerformAgentChain`:** `PerformAgentChain` runs on every subtask iteration, including iterations where the provider has not changed. Normalizing there would clear reasoning signatures unconditionally, eliminating any KV cache benefit from the LLM backend. `processChain` is only called when new user input is being stored — exactly the moment a new conversation boundary exists and normalization is safe. **Why add a continuation Human message:** When a provider is switched and the last message in the chain is not a Human message (e.g., the previous turn ended with a tool call whose response was just added), the new provider would receive a chain ending with a Tool message. Some providers reject this or behave unpredictably. Appending `"Continue from where we left off"` creates a clean conversation turn that any provider can handle correctly without changing the semantic content of the task. **Why reuse the executor on provider switch:** The Docker containers holding pentest tool environments are expensive to create and are unrelated to the LLM provider. `SetProvider` replaces only the LLM client and `tcIDTemplate` inside the existing `flowProvider` struct under a mutex, leaving the executor, image reference, and all log workers intact. **Why the restart scenario still requires user input:** This is intentional. The backend cannot safely resume execution automatically after an unknown interruption — the user should confirm the intent to continue. Once the user submits any input, `processChain` fires, fixes the chain, and execution resumes. This matches the existing graceful degradation model. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
yindo added the pull-request label 2026-06-06 22:09:57 -04:00
yindo closed this issue 2026-06-06 22:09:58 -04:00
Sign in to join this conversation.
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: vxcontrol/pentagi#267