526 lines
20 KiB
Rust
526 lines
20 KiB
Rust
//! Manages how we serialize blocklist AI data for persistence.
|
|
#![cfg_attr(not(feature = "local_fs"), allow(dead_code))]
|
|
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
use anyhow::anyhow;
|
|
use chrono::{DateTime, Local};
|
|
use serde::{Deserialize, Deserializer, Serialize};
|
|
use uuid::Uuid;
|
|
|
|
use super::AIQueryHistoryOutputStatus;
|
|
use crate::ai::agent::conversation::AIConversationId;
|
|
use crate::ai::agent::{
|
|
AIAgentActionType, AIAgentAttachment, AIAgentContext, AIAgentExchangeId, AIAgentInput,
|
|
AIAgentPtyWriteMode, AskUserQuestionItem, FileLocations, PassiveSuggestionResultType,
|
|
ReadFilesRequest, RequestComputerUseRequest, SearchCodebaseRequest, UseComputerRequest,
|
|
UserQueryMode,
|
|
};
|
|
use crate::ai::llms::LLMId;
|
|
use crate::terminal::model::block::{BlockId, SerializedBlock};
|
|
/// Data we persist for each [`AIAgentExchange`] for use in history. Does not contain output data.
|
|
#[derive(Debug, Deserialize, Clone)]
|
|
pub struct PersistedAIInput {
|
|
pub(crate) exchange_id: AIAgentExchangeId,
|
|
pub(crate) conversation_id: AIConversationId,
|
|
pub(crate) start_ts: DateTime<Local>,
|
|
pub(crate) inputs: Vec<PersistedAIInputType>,
|
|
pub(crate) output_status: AIQueryHistoryOutputStatus,
|
|
pub(crate) working_directory: Option<String>,
|
|
// TODO(CORE-3546): pub(crate) shell: Option<AvailableShell>,
|
|
pub(crate) model_id: LLMId,
|
|
#[allow(unused)]
|
|
pub(crate) coding_model_id: LLMId,
|
|
}
|
|
|
|
/// Pieces of data we need to persist for each [`AIAgentExchange`]'s input for session restoration.
|
|
///
|
|
/// Note: Only Query is actually used - it's used for up-arrow history.
|
|
/// TODO(roland): consider removing the ai_queries table and getting queries from tasks as well.
|
|
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
|
|
pub(crate) enum PersistedAIInputType {
|
|
Query {
|
|
text: String,
|
|
#[serde(default)]
|
|
context: Arc<[AIAgentContext]>,
|
|
#[serde(default)]
|
|
referenced_attachments: HashMap<String, AIAgentAttachment>,
|
|
},
|
|
}
|
|
|
|
impl TryFrom<&AIAgentInput> for PersistedAIInputType {
|
|
type Error = anyhow::Error;
|
|
|
|
fn try_from(input: &AIAgentInput) -> Result<Self, Self::Error> {
|
|
match input {
|
|
AIAgentInput::UserQuery {
|
|
query,
|
|
context,
|
|
referenced_attachments,
|
|
..
|
|
} => Ok(Self::Query {
|
|
text: query.clone(),
|
|
context: context.clone(),
|
|
referenced_attachments: referenced_attachments.clone(),
|
|
}),
|
|
AIAgentInput::AutoCodeDiffQuery { query, context } => Ok(Self::Query {
|
|
text: query.clone(),
|
|
context: context.clone(),
|
|
referenced_attachments: Default::default(),
|
|
}),
|
|
AIAgentInput::PassiveSuggestionResult { suggestion: PassiveSuggestionResultType::Prompt { prompt }, context, .. } => Ok(Self::Query {
|
|
text: prompt.clone(),
|
|
context: context.clone(),
|
|
referenced_attachments: Default::default(),
|
|
}),
|
|
AIAgentInput::PassiveSuggestionResult { suggestion: PassiveSuggestionResultType::CodeDiff { .. }, .. } => Err(anyhow!(
|
|
"PassiveSuggestionResult::CodeDiff is not persisted as a query."
|
|
)),
|
|
AIAgentInput::CommandCompletionAssessment { .. }
|
|
| AIAgentInput::ActionResult { .. }
|
|
| AIAgentInput::ResumeConversation { .. }
|
|
| AIAgentInput::InitProjectRules { .. }
|
|
| AIAgentInput::CreateEnvironment { .. }
|
|
| AIAgentInput::TriggerPassiveSuggestion { .. }
|
|
| AIAgentInput::CreateNewProject { .. }
|
|
| AIAgentInput::CloneRepository { .. }
|
|
| AIAgentInput::CodeReview { .. }
|
|
| AIAgentInput::FetchReviewComments { .. }
|
|
| AIAgentInput::SummarizeConversation { .. }
|
|
| AIAgentInput::InvokeSkill { .. }
|
|
| AIAgentInput::StartFromAmbientRunPrompt { .. }
|
|
| AIAgentInput::MessagesReceivedFromAgents { .. }
|
|
| AIAgentInput::EventsFromAgents { .. }
|
|
| AIAgentInput::OrchestrationConfigUpdate { .. } => Err(anyhow::anyhow!(
|
|
"This input type is not persisted. Only Query inputs are persisted for up-arrow history."
|
|
)),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl TryFrom<PersistedAIInputType> for AIAgentInput {
|
|
type Error = anyhow::Error;
|
|
|
|
fn try_from(value: PersistedAIInputType) -> Result<Self, Self::Error> {
|
|
match value {
|
|
PersistedAIInputType::Query {
|
|
text,
|
|
context,
|
|
referenced_attachments,
|
|
} => Ok(Self::UserQuery {
|
|
query: text,
|
|
context,
|
|
referenced_attachments,
|
|
static_query_type: None,
|
|
user_query_mode: UserQueryMode::default(),
|
|
running_command: None,
|
|
intended_agent: None,
|
|
}),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
|
pub(crate) enum PersistedPtyWriteMode {
|
|
Raw,
|
|
Line,
|
|
Block,
|
|
}
|
|
|
|
impl From<PersistedPtyWriteMode> for AIAgentPtyWriteMode {
|
|
fn from(value: PersistedPtyWriteMode) -> Self {
|
|
match value {
|
|
PersistedPtyWriteMode::Raw => Self::Raw,
|
|
PersistedPtyWriteMode::Block => Self::Block,
|
|
PersistedPtyWriteMode::Line => Self::Line,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl From<AIAgentPtyWriteMode> for PersistedPtyWriteMode {
|
|
fn from(value: AIAgentPtyWriteMode) -> Self {
|
|
match value {
|
|
AIAgentPtyWriteMode::Raw => Self::Raw,
|
|
AIAgentPtyWriteMode::Block => Self::Block,
|
|
AIAgentPtyWriteMode::Line => Self::Line,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub(crate) enum PersistedAIAgentActionType {
|
|
RequestCommandOutput {
|
|
command: String,
|
|
},
|
|
WriteToLongRunningShellCommand {
|
|
block_id: BlockId,
|
|
input: bytes::Bytes,
|
|
mode: PersistedPtyWriteMode,
|
|
},
|
|
RequestFileEdits {
|
|
file_names: Vec<String>,
|
|
},
|
|
GetFiles {
|
|
file_names: Vec<String>,
|
|
},
|
|
GetRelevantFiles {
|
|
query: String,
|
|
partial_paths: Option<Vec<String>>,
|
|
codebase_path: Option<String>,
|
|
},
|
|
Grep {
|
|
queries: Vec<String>,
|
|
path: String,
|
|
},
|
|
FileGlob {
|
|
patterns: Vec<String>,
|
|
path: Option<String>,
|
|
},
|
|
FileGlobV2 {
|
|
patterns: Vec<String>,
|
|
search_dir: Option<String>,
|
|
},
|
|
ReadMCPResource {
|
|
server_id: Option<Uuid>,
|
|
name: String,
|
|
uri: Option<String>,
|
|
},
|
|
CallMCPTool {
|
|
server_id: Option<Uuid>,
|
|
name: String,
|
|
input: serde_json::Value,
|
|
},
|
|
SuggestNewConversation {
|
|
message_id: String,
|
|
},
|
|
SuggestPrompt,
|
|
OpenCodeReview,
|
|
InitProject,
|
|
UseComputer {
|
|
action_summary: String,
|
|
#[serde(deserialize_with = "deserialize_targeted_actions")]
|
|
actions: Vec<computer_use::TargetedAction>,
|
|
screenshot_params: Option<computer_use::ScreenshotParams>,
|
|
},
|
|
RequestComputerUse {
|
|
task_summary: String,
|
|
screenshot_params: Option<computer_use::ScreenshotParams>,
|
|
},
|
|
AskUserQuestion {
|
|
questions: Vec<AskUserQuestionItem>,
|
|
},
|
|
|
|
FetchConversation {
|
|
conversation_id: String,
|
|
},
|
|
|
|
/// Actions that don't need data persisted (since they're restored from conversation tasks) can be mapped to this.
|
|
NotPersisted,
|
|
}
|
|
|
|
/// Deserializes the persisted `UseComputer` actions, accepting both the current `{ action, target }`
|
|
/// shape and the legacy bare-`Action` shape.
|
|
///
|
|
/// Conversations persisted before `actions` became `Vec<TargetedAction>` stored each element as a
|
|
/// bare [`computer_use::Action`]; those decode with the target defaulting to `Target::Screen`. New
|
|
/// data round-trips unchanged, and serialization still emits the `{ action, target }` shape via the
|
|
/// derived `Serialize` impl.
|
|
fn deserialize_targeted_actions<'de, D>(
|
|
deserializer: D,
|
|
) -> Result<Vec<computer_use::TargetedAction>, D::Error>
|
|
where
|
|
D: Deserializer<'de>,
|
|
{
|
|
// Accepts either the new `{ action, target }` wrapper (with `target` optional) or a bare legacy
|
|
// `Action` value. `Action`'s variant names never collide with the `action`/`target` keys, so
|
|
// the untagged match is unambiguous.
|
|
#[derive(Deserialize)]
|
|
#[serde(untagged)]
|
|
enum TargetedActionCompat {
|
|
Targeted {
|
|
action: computer_use::Action,
|
|
#[serde(default)]
|
|
target: computer_use::Target,
|
|
},
|
|
Bare(computer_use::Action),
|
|
}
|
|
|
|
let actions = Vec::<TargetedActionCompat>::deserialize(deserializer)?;
|
|
Ok(actions
|
|
.into_iter()
|
|
.map(|compat| match compat {
|
|
TargetedActionCompat::Targeted { action, target } => {
|
|
computer_use::TargetedAction { action, target }
|
|
}
|
|
TargetedActionCompat::Bare(action) => computer_use::TargetedAction::screen(action),
|
|
})
|
|
.collect())
|
|
}
|
|
|
|
impl From<&AIAgentActionType> for PersistedAIAgentActionType {
|
|
fn from(value: &AIAgentActionType) -> Self {
|
|
match value {
|
|
AIAgentActionType::RequestCommandOutput { command, .. } => Self::RequestCommandOutput {
|
|
command: command.clone(),
|
|
},
|
|
AIAgentActionType::WriteToLongRunningShellCommand {
|
|
block_id,
|
|
input,
|
|
mode,
|
|
} => Self::WriteToLongRunningShellCommand {
|
|
block_id: block_id.clone(),
|
|
input: input.clone(),
|
|
mode: (*mode).into(),
|
|
},
|
|
AIAgentActionType::RequestFileEdits { file_edits, .. } => Self::RequestFileEdits {
|
|
file_names: file_edits
|
|
.iter()
|
|
.filter_map(|diff| diff.file())
|
|
.map(ToOwned::to_owned)
|
|
.collect(),
|
|
},
|
|
AIAgentActionType::ReadFiles(ReadFilesRequest { locations: files }) => Self::GetFiles {
|
|
file_names: files.iter().map(|f| f.name.clone()).collect(),
|
|
},
|
|
AIAgentActionType::SearchCodebase(SearchCodebaseRequest {
|
|
query,
|
|
partial_paths,
|
|
codebase_path,
|
|
}) => Self::GetRelevantFiles {
|
|
query: query.clone(),
|
|
partial_paths: partial_paths.clone(),
|
|
codebase_path: codebase_path.clone(),
|
|
},
|
|
AIAgentActionType::Grep { queries, path } => Self::Grep {
|
|
queries: queries.clone(),
|
|
path: path.clone(),
|
|
},
|
|
AIAgentActionType::FileGlob { patterns, path } => Self::FileGlob {
|
|
patterns: patterns.clone(),
|
|
path: path.clone(),
|
|
},
|
|
AIAgentActionType::FileGlobV2 {
|
|
patterns,
|
|
search_dir,
|
|
} => Self::FileGlobV2 {
|
|
patterns: patterns.clone(),
|
|
search_dir: search_dir.clone(),
|
|
},
|
|
AIAgentActionType::CallMCPTool {
|
|
server_id,
|
|
name,
|
|
input,
|
|
} => Self::CallMCPTool {
|
|
server_id: *server_id,
|
|
name: name.clone(),
|
|
input: input.clone(),
|
|
},
|
|
AIAgentActionType::ReadMCPResource {
|
|
server_id,
|
|
name,
|
|
uri,
|
|
} => Self::ReadMCPResource {
|
|
server_id: *server_id,
|
|
name: name.clone(),
|
|
uri: uri.clone(),
|
|
},
|
|
AIAgentActionType::SuggestNewConversation { message_id } => {
|
|
Self::SuggestNewConversation {
|
|
message_id: message_id.clone(),
|
|
}
|
|
}
|
|
AIAgentActionType::SuggestPrompt { .. } => Self::SuggestPrompt,
|
|
AIAgentActionType::OpenCodeReview => Self::OpenCodeReview,
|
|
AIAgentActionType::InsertCodeReviewComments { .. } => Self::NotPersisted,
|
|
AIAgentActionType::InitProject => Self::InitProject,
|
|
AIAgentActionType::ReadDocuments(_)
|
|
| AIAgentActionType::EditDocuments(_)
|
|
| AIAgentActionType::CreateDocuments(_)
|
|
| AIAgentActionType::ReadShellCommandOutput { .. }
|
|
| AIAgentActionType::ReadSkill(_)
|
|
| AIAgentActionType::UploadArtifact(_)
|
|
| AIAgentActionType::TransferShellCommandControlToUser { .. } => Self::NotPersisted,
|
|
AIAgentActionType::UseComputer(req) => Self::UseComputer {
|
|
action_summary: req.action_summary.clone(),
|
|
actions: req.actions.clone(),
|
|
screenshot_params: req.screenshot_params,
|
|
},
|
|
AIAgentActionType::RequestComputerUse(req) => Self::RequestComputerUse {
|
|
task_summary: req.task_summary.clone(),
|
|
screenshot_params: req.screenshot_params,
|
|
},
|
|
AIAgentActionType::AskUserQuestion { questions } => Self::AskUserQuestion {
|
|
questions: questions.clone(),
|
|
},
|
|
AIAgentActionType::FetchConversation { conversation_id } => Self::FetchConversation {
|
|
conversation_id: conversation_id.clone(),
|
|
},
|
|
AIAgentActionType::StartAgent { .. } => Self::NotPersisted,
|
|
AIAgentActionType::SendMessageToAgent { .. } => Self::NotPersisted,
|
|
// Orchestrate is rendered from the in-history tool call message;
|
|
// there is no per-action state we need to persist locally.
|
|
AIAgentActionType::RunAgents(_) => Self::NotPersisted,
|
|
// The wait is dropped on restart; the unresolved tool call
|
|
// stays in the transcript as an orphan until the next
|
|
// outbound request triggers the server's supersede.
|
|
AIAgentActionType::WaitForEvents { .. } => Self::NotPersisted,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl TryFrom<PersistedAIAgentActionType> for AIAgentActionType {
|
|
type Error = anyhow::Error;
|
|
|
|
fn try_from(value: PersistedAIAgentActionType) -> Result<Self, Self::Error> {
|
|
match value {
|
|
PersistedAIAgentActionType::RequestCommandOutput { command, .. } => {
|
|
Ok(Self::RequestCommandOutput {
|
|
command,
|
|
rationale: None,
|
|
is_read_only: None,
|
|
is_risky: None,
|
|
uses_pager: None,
|
|
// TODO(zachbai): Support restoring this value from persisted type.
|
|
wait_until_completion: false,
|
|
citations: vec![],
|
|
})
|
|
}
|
|
PersistedAIAgentActionType::WriteToLongRunningShellCommand {
|
|
block_id,
|
|
input,
|
|
mode,
|
|
} => Ok(Self::WriteToLongRunningShellCommand {
|
|
block_id: block_id.clone(),
|
|
input: input.clone(),
|
|
mode: mode.into(),
|
|
}),
|
|
PersistedAIAgentActionType::GetRelevantFiles {
|
|
query,
|
|
partial_paths,
|
|
codebase_path,
|
|
} => Ok(Self::SearchCodebase(SearchCodebaseRequest {
|
|
query,
|
|
partial_paths,
|
|
codebase_path,
|
|
})),
|
|
PersistedAIAgentActionType::RequestFileEdits { .. } => {
|
|
// TODO(CODE-301): Implement proper restoration for suggested diffs.
|
|
//
|
|
// The current "implementation" is incomplete and does not actually persist any
|
|
// diff content. For now, we just ignore the suggested diff actions altogether,
|
|
// instead of restoring diffs with no content.
|
|
Err(anyhow!("Restoration for RequestFileEdits is unsupported. "))
|
|
}
|
|
PersistedAIAgentActionType::GetFiles { file_names } => {
|
|
Ok(Self::ReadFiles(ReadFilesRequest {
|
|
locations: file_names
|
|
.into_iter()
|
|
.map(|name| FileLocations {
|
|
name,
|
|
lines: Vec::new(),
|
|
})
|
|
.collect(),
|
|
}))
|
|
}
|
|
PersistedAIAgentActionType::Grep { queries, path } => Ok(Self::Grep { queries, path }),
|
|
PersistedAIAgentActionType::FileGlob { patterns, path } => {
|
|
Ok(Self::FileGlob { patterns, path })
|
|
}
|
|
PersistedAIAgentActionType::FileGlobV2 {
|
|
patterns,
|
|
search_dir,
|
|
} => Ok(Self::FileGlobV2 {
|
|
patterns,
|
|
search_dir,
|
|
}),
|
|
PersistedAIAgentActionType::CallMCPTool {
|
|
server_id,
|
|
name,
|
|
input,
|
|
} => Ok(Self::CallMCPTool {
|
|
server_id,
|
|
name,
|
|
input,
|
|
}),
|
|
PersistedAIAgentActionType::ReadMCPResource {
|
|
server_id,
|
|
name,
|
|
uri,
|
|
} => Ok(Self::ReadMCPResource {
|
|
server_id,
|
|
name,
|
|
uri,
|
|
}),
|
|
PersistedAIAgentActionType::SuggestNewConversation { message_id } => {
|
|
Ok(Self::SuggestNewConversation {
|
|
message_id: message_id.clone(),
|
|
})
|
|
}
|
|
PersistedAIAgentActionType::SuggestPrompt => {
|
|
Err(anyhow!("Restoration for suggested prompts is unsupported."))
|
|
}
|
|
PersistedAIAgentActionType::OpenCodeReview => Ok(Self::OpenCodeReview),
|
|
PersistedAIAgentActionType::InitProject => Ok(Self::InitProject),
|
|
PersistedAIAgentActionType::UseComputer {
|
|
action_summary,
|
|
actions,
|
|
screenshot_params,
|
|
} => Ok(Self::UseComputer(UseComputerRequest {
|
|
action_summary,
|
|
actions,
|
|
screenshot_params,
|
|
})),
|
|
PersistedAIAgentActionType::RequestComputerUse {
|
|
task_summary,
|
|
screenshot_params,
|
|
} => Ok(Self::RequestComputerUse(RequestComputerUseRequest {
|
|
task_summary,
|
|
screenshot_params,
|
|
})),
|
|
PersistedAIAgentActionType::AskUserQuestion { questions } => {
|
|
Ok(Self::AskUserQuestion { questions })
|
|
}
|
|
PersistedAIAgentActionType::FetchConversation { conversation_id } => {
|
|
Ok(Self::FetchConversation { conversation_id })
|
|
}
|
|
PersistedAIAgentActionType::NotPersisted => Err(anyhow!(
|
|
"Restoration is handled through conversation tasks, not persisted blocks."
|
|
)),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// The types of "blocks" we can store in our SQLite database for session restoration. Only command
|
|
/// blocks are true [`crate::terminal::model::block::Block`]s.
|
|
///
|
|
/// TODO(roland): now that there is no AI serialized block, consider removing this enum wrapper
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub enum SerializedBlockListItem {
|
|
Command { block: Box<SerializedBlock> },
|
|
}
|
|
|
|
impl SerializedBlockListItem {
|
|
pub(crate) fn start_ts(&self) -> Option<DateTime<Local>> {
|
|
match self {
|
|
Self::Command { block } => block.start_ts,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl From<crate::persistence::model::Block> for SerializedBlockListItem {
|
|
fn from(value: crate::persistence::model::Block) -> Self {
|
|
Self::Command {
|
|
block: Box::new(SerializedBlock::from(value)),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl From<SerializedBlock> for SerializedBlockListItem {
|
|
fn from(value: SerializedBlock) -> Self {
|
|
Self::Command {
|
|
block: Box::new(value),
|
|
}
|
|
}
|
|
}
|