refactor(runtime): use StakAI message models#732
Open
ahmedhesham6 wants to merge 5 commits into
Open
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
Refactors the CLI/runtime message flow to use StakAI-native Message, ToolCall, and ContentPart models end-to-end, removing OpenAI-shaped bridge types and updating storage/checkpoint and UI consumers accordingly.
Changes:
- Introduces shared
agent_runtimemodel types inlibs/shared(tool result/progress + ask_user types) and updates TUI/CLI/MCP/server to consume them. - Switches API client/provider execution to StakAI
GenerateRequest/StreamEvent/Usage, and removes the server-side message bridge. - Updates session/storage and rendering paths to work with StakAI message parts (tool calls/results, images) directly.
Reviewed changes
Copilot reviewed 70 out of 71 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tui/src/services/textarea.rs | Test refactor using if let ... && let ... chaining. |
| tui/src/services/message.rs | Updates tool-call handling to StakAI ToolCall shape and JSON-value arguments. |
| tui/src/services/markdown_renderer.rs | Minor test assertion simplifications. |
| tui/src/services/image_upload.rs | Switches image parts to StakAI ContentPart constructors. |
| tui/src/services/handlers/tool.rs | Updates tool handling to StakAI ToolCall + shared runtime result/progress types. |
| tui/src/services/handlers/shell.rs | Updates shell tool-call parsing and tool-call/result construction for StakAI types. |
| tui/src/services/handlers/mod.rs | Adjusts approval lookups now that ToolCall equality semantics changed. |
| tui/src/services/handlers/input.rs | Same approval lookup update as handlers/mod.rs. |
| tui/src/services/handlers/dialog.rs | Migrates dialog tool-call handling to StakAI ToolCall. |
| tui/src/services/handlers/ask_user.rs | Migrates ask_user handling to shared runtime ask_user types + StakAI ToolCall. |
| tui/src/services/file_diff.rs | Reads diff args from StakAI ToolCall.arguments JSON value. |
| tui/src/services/detect_term.rs | Minor iterator change in tests (next_back). |
| tui/src/services/changeset.rs | Migrates changeset tracking/reverse edits to StakAI ToolCall. |
| tui/src/services/board_tasks.rs | Updates message scanning logic to handle new argument/result representations. |
| tui/src/services/bash_block.rs | Migrates tool rendering paths to StakAI ToolCall + shared runtime types. |
| tui/src/services/auto_approve.rs | Migrates auto-approve policy evaluation to StakAI ToolCall arguments. |
| tui/src/services/approval_bar.rs | Migrates approval-bar tool labeling/parsing to StakAI ToolCall. |
| tui/src/event_loop.rs | Updates tool name extraction to StakAI ToolCall shape. |
| tui/src/app/types.rs | Switches app state to StakAI ContentPart/ToolCall and shared runtime types. |
| tui/src/app/events.rs | Updates TUI event payloads to StakAI messages/tool calls/content parts. |
| libs/shared/src/remote_connection.rs | Uses shared runtime ToolCallResultProgress for MCP progress notifications. |
| libs/shared/src/models/mod.rs | Exposes new agent_runtime module; removes stakai_adapter. |
| libs/shared/src/models/llm.rs | Removes legacy LLM input/message/tool/delta types no longer used after refactor. |
| libs/shared/src/models/integrations/mcp.rs | Updates MCP integration to shared runtime ToolCallResultStatus; removes chat-message status helper. |
| libs/shared/src/models/async_manifest.rs | Stores pending tool calls from StakAI ToolCall directly. |
| libs/shared/src/models/agent_runtime.rs | Adds shared runtime types for tool results/progress/task updates + re-exports ask_user types. |
| libs/server/src/session_actor.rs | Removes storage bridge usage; persists Vec<stakai::Message> directly. |
| libs/server/src/routes.rs | Returns checkpoint messages without bridging conversions. |
| libs/server/src/message_bridge.rs | Removes obsolete ChatMessage ↔ StakAI message adapter (deleted). |
| libs/server/src/lib.rs | Stops exporting message_bridge. |
| libs/server/README.md | Updates docs to reflect StakAI-native storage edge. |
| libs/mcp/server/src/local_tools.rs | Switches tool progress/task updates imports to shared runtime types. |
| libs/mcp/client/src/local.rs | Updates progress type import to shared runtime model. |
| libs/mcp/client/src/lib.rs | Updates progress type import to shared runtime model. |
| libs/api/src/storage.rs | Checkpoint/session storage models now store Vec<stakai::Message>. |
| libs/api/src/stakpak/models.rs | Same storage model update for stakpak backend models. |
| libs/api/src/models.rs | Refactors API models to StakAI messages/tools/requests; adds stream event/usage conversion. |
| libs/api/src/local/tests.rs | Updates local-storage tests to StakAI message parts (tool results, trimming checks). |
| libs/api/src/local/hooks/task_board_context/mod.rs | Builds StakAI GenerateRequest for hook-driven inference. |
| libs/api/src/local/hooks/inline_scratchpad_context/mod.rs | Same: StakAI request construction for scratchpad hook. |
| libs/api/src/local/hooks/file_scratchpad_context/mod.rs | Same: StakAI request construction for file scratchpad hook. |
| libs/api/src/local/context_managers/simple_context_manager.rs | Context manager now operates on stakai::Message. |
| libs/api/src/local/context_managers/scratchpad_context_manager.rs | Context manager now operates on stakai::Message. |
| libs/api/src/local/context_managers/mod.rs | Trait signature updated to StakAI messages. |
| libs/api/src/local/context_managers/file_scratchpad_context_manager.rs | Reworks history/tool-call reconstruction around StakAI message parts. |
| libs/api/src/local/context_managers/common.rs | Shared history conversion updated for StakAI tool-call/tool-result parts and role labels. |
| libs/api/src/lib.rs | AgentProvider trait now returns StakAI completion/stream events. |
| libs/api/src/client/stakai.rs | New StakAI inference client + provider registry builder. |
| libs/api/src/client/provider.rs | Provider now executes StakAI generate/stream, emits AgentStreamEvent, and builds StakAI messages. |
| libs/api/src/client/mod.rs | Wires in new client::stakai module and removes shared adapter client. |
| libs/ai/src/providers/gemini/convert.rs | Minor test comment update after adapter removal. |
| libs/ai/docs/PROVIDER_OPTIONS_PLAN.md | Updates docs to remove references to the deleted adapter layer. |
| cli/src/utils/agent_context.rs | Adds helper to strip injected context blocks from user-visible text. |
| cli/src/commands/sessions/tests.rs | Updates session command tests to StakAI message/content-part shapes. |
| cli/src/commands/sessions/output.rs | Renders tool calls/results by iterating StakAI message parts. |
| cli/src/commands/sessions/messages.rs | Filters messages using StakAI roles/messages; adjusts system-role behavior. |
| cli/src/commands/agent/run/tui.rs | Sends StakAI tool calls to TUI. |
| cli/src/commands/agent/run/tooling.rs | Runs tools using StakAI tool-call names/arguments JSON values. |
| cli/src/commands/agent/run/stream.rs | Replaces OpenAI-delta accumulator with StakAI StreamEvent accumulator and new completion shape. |
| cli/src/commands/agent/run/renderer.rs | Renders final output from StakAI messages. |
| cli/src/commands/agent/run/pause.rs | Detects pending tool calls/results using StakAI message parts. |
| cli/src/commands/agent/run/mode_interactive.rs | Migrates interactive loop to StakAI messages/tool calls; updates tool/result handling. |
| cli/src/commands/agent/run/mode_async.rs | Migrates async mode to StakAI messages/tool calls; updates usage conversion and pause flow. |
| cli/src/commands/agent/run/mcp_init.rs | Uses shared runtime progress + StakAI Tool list. |
| cli/src/commands/agent/run/helpers.rs | Migrates helper constructors (user/system/tool result) to StakAI messages/parts. |
| cli/src/commands/agent/run/checkpoint.rs | Migrates checkpoint replay to StakAI messages/parts and strips injected context from displayed user text. |
| cli/src/commands/acp/server.rs | Migrates ACP server agent to StakAI stream/events/messages; updates tool execution plumbing. |
| cli/src/commands/acp/fs_handler.rs | Migrates ACP FS tool execution to StakAI tool-call arguments JSON values. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
This commit refactors how tool results are extracted and handled in
several areas:
- `cli/src/commands/agent/run/checkpoint.rs`:
- Introduces constants for known tool result markers (e.g.,
`TOOL_CALL_CANCELLED`, `TOOL_CALL_REJECTED`).
- Implements `tool_result_status_from_content` to correctly infer the
`ToolCallResultStatus` from various content formats, including
explicit markers and structured error objects.
- Enhances `tool_result_content_string` to handle non-string JSON
content more robustly.
- Adds tests to verify the extraction of cancelled and various error
statuses for tool results.
- `cli/src/commands/agent/run/tooling.rs`:
- Modifies `run_tool_call` to use a new `mcp_tool_arguments` function
for parsing tool call arguments.
- `mcp_tool_arguments` now returns a `CallToolResult` error if the
arguments are not a JSON object or null, improving argument
validation.
- Adds a helper `json_type_name` for better error messages.
- Includes a test to ensure
# Conflicts: # cli/src/commands/agent/run/mcp_init.rs # libs/shared/src/models/stakai_adapter.rs # tui/src/event_loop.rs
Comment on lines
44
to
55
| @@ -40,71 +47,65 @@ pub async fn extract_checkpoint_messages_and_tool_calls( | |||
| .find(|message| message.role != Role::User && message.role != Role::Tool) | |||
| && last_message.role == Role::Assistant | |||
| { | |||
| last_message.content = Some(MessageContent::String(format!( | |||
| last_message.content = MessageContent::Text(format!( | |||
| "{}\n<checkpoint_id>{}</checkpoint_id>", | |||
| last_message | |||
| .content | |||
| .as_ref() | |||
| .unwrap_or(&MessageContent::String(String::new())), | |||
| last_message.content.text().unwrap_or_default(), | |||
| checkpoint_id | |||
| ))); | |||
| )); | |||
| } | |||
Comment on lines
+944
to
951
| usage: final_usage, .. | ||
| } => { | ||
| usage = final_usage.clone(); | ||
| } | ||
| StreamEvent::Start { .. } | ||
| | StreamEvent::ReasoningDelta { .. } | ||
| | StreamEvent::Error { .. } => {} | ||
| } |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Refactors the CLI/runtime message path to use StakAI-native message and tool-call models instead of OpenAI-shaped bridge types.
Related Issues
None.
Changes Made
Message,ToolCall, andContentPartshapes directly.Testing
rustfmt --edition 2024 --checkon staged Rust filescargo clippy -p stakpak -p stakpak-tui -p stakai -p stakpak-api -p stakpak-shared -p stakpak-mcp-client -p stakpak-mcp-server -p stakpak-server --all-targets -- -D warningsScreenshots (if applicable)
N/A
Breaking Changes
None expected.