use std::sync::Arc;
use ag_agent::{self as agent, AgentKind, AgentModel, ReasoningLevel, SessionStats, SpeedMode};
use ag_session::{FocusedReviewStatus, SessionMessageKind};
use async_trait::async_trait;
use sqlx::SqlitePool;
use tracing::warn;
use super::review::SessionReviewRequestRow;
use super::session_message::SessionMessageStore;
use super::session_snapshot::SessionSnapshotStore;
use super::status;
use crate::DbError;
use crate::timestamp::TimestampSource;
pub struct SessionTurnMetadata {
pub applied_personality_id: Option<String>,
pub applied_personality_prompt_hash: Option<String>,
pub instruction_conversation_id: Option<String>,
pub model: String,
pub provider_conversation_id: Option<String>,
pub questions_json: String,
pub summary: String,
pub token_usage_delta: SessionStats,
}
pub struct PersistedSessionCreation<'a> {
pub agent: &'a str,
pub base_branch: &'a str,
pub id: &'a str,
pub is_draft: bool,
pub model: &'a str,
pub orchestration_task_id: Option<i64>,
pub parent_session_id: Option<&'a str>,
pub personality_id: Option<&'a str>,
pub project_id: i64,
pub reasoning_level: ReasoningLevel,
pub role: Option<&'a str>,
pub speed_mode: SpeedMode,
pub status: &'a str,
}
pub struct ForkSessionSnapshot<'a> {
pub new_session_id: &'a str,
pub source_session_id: &'a str,
pub status: &'a str,
}
pub struct SessionRow {
pub added_lines: i64,
pub agent: String,
pub base_branch: String,
pub created_at: i64,
pub deleted_lines: i64,
pub has_diff: Option<bool>,
pub id: String,
pub in_progress_started_at: Option<i64>,
pub in_progress_total_seconds: i64,
pub input_tokens: i64,
pub is_draft: bool,
pub model: String,
pub output_tokens: i64,
pub parent_session_id: Option<String>,
pub personality_id: Option<String>,
pub project_id: Option<i64>,
pub prompt: String,
pub published_upstream_ref: Option<String>,
pub questions: Option<String>,
pub reasoning_level_override: Option<String>,
pub review_request: Option<SessionReviewRequestRow>,
pub role: Option<String>,
pub size: String,
pub speed_mode: String,
pub status: String,
pub summary: Option<String>,
pub title: Option<String>,
pub updated_at: i64,
}
pub struct SessionListRow {
pub added_lines: i64,
pub agent: String,
pub base_branch: String,
pub created_at: i64,
pub deleted_lines: i64,
pub has_diff: Option<bool>,
pub id: String,
pub in_progress_started_at: Option<i64>,
pub in_progress_total_seconds: i64,
pub input_tokens: i64,
pub is_draft: bool,
pub model: String,
pub output_tokens: i64,
pub parent_session_id: Option<String>,
pub personality_id: Option<String>,
pub project_id: Option<i64>,
pub published_upstream_ref: Option<String>,
pub reasoning_level_override: Option<String>,
pub review_request: Option<SessionReviewRequestRow>,
pub role: Option<String>,
pub size: String,
pub speed_mode: String,
pub status: String,
pub title: Option<String>,
pub updated_at: i64,
}
#[derive(sqlx::FromRow)]
pub struct SessionAgentModelRow {
pub agent: String,
pub id: String,
pub model: String,
pub status: String,
}
pub struct SessionDetailRow {
pub prompt: String,
pub questions: Option<String>,
pub summary: Option<String>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SessionMessageRow {
pub content: String,
pub kind: String,
pub position: i64,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SessionFocusedReviewRow {
pub diff_hash: String,
pub session_id: String,
pub text: String,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SessionPersonalityState {
pub applied_personality_id: Option<String>,
pub applied_personality_prompt_hash: Option<String>,
pub personality_id: Option<String>,
}
#[async_trait]
pub trait SessionRepository: Send + Sync {
async fn append_session_message(
&self,
id: &str,
kind: SessionMessageKind,
content: &str,
) -> Result<(), DbError>;
async fn backfill_session_project(&self, project_id: i64) -> Result<(), DbError>;
async fn delete_session(&self, id: &str) -> Result<(), DbError>;
async fn get_session_base_branch(&self, id: &str) -> Result<Option<String>, DbError>;
async fn get_session_parent_session_id(&self, id: &str) -> Result<Option<String>, DbError>;
async fn get_session_stack_base_commit_hash(&self, id: &str)
-> Result<Option<String>, DbError>;
async fn get_session_instruction_conversation_id(
&self,
id: &str,
) -> Result<Option<String>, DbError>;
async fn get_session_provider_conversation_id(
&self,
id: &str,
) -> Result<Option<String>, DbError>;
async fn insert_draft_session(
&self,
id: &str,
model: &str,
base_branch: &str,
status: &str,
project_id: i64,
) -> Result<(), DbError>;
async fn insert_stacked_draft_session(
&self,
id: &str,
model: &str,
base_branch: &str,
status: &str,
parent_session_id: &str,
project_id: i64,
) -> Result<(), DbError>;
async fn insert_session(
&self,
id: &str,
model: &str,
base_branch: &str,
status: &str,
project_id: i64,
) -> Result<(), DbError>;
async fn insert_session_with_agent(
&self,
session: PersistedSessionCreation<'_>,
) -> Result<(), DbError>;
async fn fork_session_snapshot(&self, snapshot: ForkSessionSnapshot<'_>)
-> Result<(), DbError>;
async fn load_session(&self, session_id: &str) -> Result<Option<SessionRow>, DbError>;
async fn load_active_session_agent_models(&self) -> Result<Vec<SessionAgentModelRow>, DbError>;
#[cfg(any(test, feature = "test-utils"))]
async fn load_sessions(&self) -> Result<Vec<SessionRow>, DbError>;
async fn load_sessions_for_project(
&self,
project_id: i64,
) -> Result<Vec<SessionListRow>, DbError>;
async fn load_session_detail(
&self,
session_id: &str,
) -> Result<Option<SessionDetailRow>, DbError>;
async fn load_session_messages(
&self,
session_id: &str,
) -> Result<Vec<SessionMessageRow>, DbError>;
async fn load_session_focused_reviews_for_project(
&self,
project_id: i64,
) -> Result<Vec<SessionFocusedReviewRow>, DbError>;
async fn load_sessions_metadata(&self) -> Result<(i64, i64), DbError>;
async fn load_session_project_id(&self, session_id: &str) -> Result<Option<i64>, DbError>;
async fn load_session_personality_state(
&self,
session_id: &str,
) -> Result<Option<SessionPersonalityState>, DbError>;
async fn load_pending_stack_restack_session_ids(
&self,
project_id: i64,
) -> Result<Vec<String>, DbError>;
async fn load_session_published_upstream_ref(
&self,
id: &str,
) -> Result<Option<String>, DbError>;
async fn load_session_merged_commit_hash(
&self,
session_id: &str,
) -> Result<Option<String>, DbError>;
async fn load_session_archived_diff(&self, session_id: &str)
-> Result<Option<String>, DbError>;
async fn restack_child_sessions_after_parent_merge(
&self,
parent_session_id: &str,
base_branch: &str,
parent_commit_hash: Option<String>,
) -> Result<Vec<String>, DbError>;
async fn load_session_reasoning_level(
&self,
session_id: &str,
) -> Result<ReasoningLevel, DbError>;
async fn load_session_speed_mode(&self, session_id: &str) -> Result<SpeedMode, DbError>;
async fn load_session_summary(&self, session_id: &str) -> Result<Option<String>, DbError>;
async fn load_session_timestamps(
&self,
session_id: &str,
) -> Result<Option<(i64, i64)>, DbError>;
async fn persist_session_turn_metadata(
&self,
session_id: &str,
turn_metadata: &SessionTurnMetadata,
) -> Result<(), DbError>;
async fn mark_session_diff_unknown(&self, id: &str) -> Result<(), DbError>;
async fn update_session_diff_stats(
&self,
added_lines: u64,
deleted_lines: u64,
has_diff: bool,
id: &str,
size: &str,
) -> Result<(), DbError>;
async fn update_session_instruction_conversation_id(
&self,
id: &str,
provider_conversation_id: Option<String>,
) -> Result<(), DbError>;
async fn update_session_model(&self, id: &str, model: &str) -> Result<(), DbError>;
async fn update_session_personality_id(
&self,
id: &str,
personality_id: Option<String>,
) -> Result<(), DbError>;
async fn update_session_agent_model(
&self,
id: &str,
agent: &str,
model: &str,
) -> Result<(), DbError>;
async fn update_active_session_agent_model(
&self,
id: &str,
agent: &str,
model: &str,
) -> Result<(), DbError>;
async fn clear_session_draft_flag(&self, id: &str) -> Result<(), DbError>;
async fn update_session_merged_commit_hash(
&self,
id: &str,
merged_commit_hash: Option<String>,
) -> Result<(), DbError>;
async fn update_session_archived_diff(
&self,
id: &str,
archived_diff: Option<String>,
) -> Result<(), DbError>;
async fn update_session_stack_base_commit_hash(
&self,
id: &str,
stack_base_commit_hash: Option<String>,
) -> Result<(), DbError>;
async fn update_session_prompt(&self, id: &str, prompt: &str) -> Result<(), DbError>;
async fn update_session_provider_conversation_id(
&self,
id: &str,
provider_conversation_id: Option<String>,
) -> Result<(), DbError>;
async fn update_session_questions(&self, id: &str, questions: &str) -> Result<(), DbError>;
async fn update_session_reasoning_level(
&self,
id: &str,
reasoning_level: ReasoningLevel,
) -> Result<(), DbError>;
async fn update_session_speed_mode(
&self,
id: &str,
speed_mode: SpeedMode,
) -> Result<(), DbError>;
async fn update_session_published_upstream_ref(
&self,
id: &str,
published_upstream_ref: Option<String>,
) -> Result<(), DbError>;
async fn update_session_stats(&self, id: &str, stats: &SessionStats) -> Result<(), DbError>;
async fn update_session_status_with_timing_at(
&self,
id: &str,
status: &str,
timestamp_seconds: i64,
) -> Result<(), DbError>;
async fn update_session_summary(&self, id: &str, summary: &str) -> Result<(), DbError>;
async fn update_session_focused_review(
&self,
id: &str,
status: Option<FocusedReviewStatus>,
diff_hash: Option<String>,
text: Option<String>,
) -> Result<(), DbError>;
async fn update_session_title(&self, id: &str, title: &str) -> Result<(), DbError>;
async fn update_session_provisional_title(&self, id: &str, title: &str) -> Result<(), DbError>;
async fn begin_session_title_generation(
&self,
id: &str,
requires_provisional_title: bool,
) -> Result<Option<i64>, DbError>;
async fn update_session_title_for_generation(
&self,
id: &str,
expected_generation: i64,
title: &str,
) -> Result<bool, DbError>;
#[cfg(any(test, feature = "test-utils"))]
async fn update_session_created_at(&self, id: &str, created_at: i64) -> Result<(), DbError>;
#[cfg(any(test, feature = "test-utils"))]
async fn update_session_updated_at(&self, id: &str, updated_at: i64) -> Result<(), DbError>;
}
#[derive(Clone)]
pub(crate) struct SqliteSessionRepository(
SqlitePool,
Arc<dyn TimestampSource>,
SessionMessageStore,
SessionSnapshotStore,
);
impl SqliteSessionRepository {
pub(crate) fn new(pool: SqlitePool, timestamp_source: Arc<dyn TimestampSource>) -> Self {
Self(
pool.clone(),
Arc::clone(×tamp_source),
SessionMessageStore::new(pool.clone(), Arc::clone(×tamp_source)),
SessionSnapshotStore::new(pool, timestamp_source),
)
}
fn now(&self) -> i64 {
self.1.now_timestamp_seconds()
}
}
struct RequiredStringValueRow {
value: String,
}
struct SessionStatsMetadataRow {
max_updated_at: i64,
session_count: i64,
}
struct OptionalI64ValueRow {
value: Option<i64>,
}
struct SessionInstructionStateRow {
app_server_instruction_provider_conversation_id: Option<String>,
}
impl SessionInstructionStateRow {
fn into_instruction_conversation_id(self) -> Option<String> {
agent::normalize_instruction_conversation_id(
self.app_server_instruction_provider_conversation_id
.as_deref(),
)
}
}
struct SessionTimestampsRow {
created_at: i64,
updated_at: i64,
}
struct SessionRowMetadata {
added_lines: i64,
agent: String,
base_branch: String,
created_at: i64,
deleted_lines: i64,
has_diff: Option<bool>,
id: String,
in_progress_started_at: Option<i64>,
in_progress_total_seconds: i64,
input_tokens: i64,
is_draft: bool,
model: String,
output_tokens: i64,
parent_session_id: Option<String>,
personality_id: Option<String>,
project_id: Option<i64>,
published_upstream_ref: Option<String>,
reasoning_level_override: Option<String>,
role: Option<String>,
size: String,
speed_mode: String,
status: String,
title: Option<String>,
updated_at: i64,
}
impl SessionRowMetadata {
fn into_session_row(
self,
prompt: String,
questions: Option<String>,
summary: Option<String>,
review_request: Option<SessionReviewRequestRow>,
) -> SessionRow {
SessionRow {
added_lines: self.added_lines,
agent: self.agent,
base_branch: self.base_branch,
created_at: self.created_at,
deleted_lines: self.deleted_lines,
has_diff: self.has_diff,
id: self.id,
in_progress_started_at: self.in_progress_started_at,
in_progress_total_seconds: self.in_progress_total_seconds,
input_tokens: self.input_tokens,
is_draft: self.is_draft,
model: self.model,
output_tokens: self.output_tokens,
parent_session_id: self.parent_session_id,
personality_id: self.personality_id,
project_id: self.project_id,
prompt,
published_upstream_ref: self.published_upstream_ref,
questions,
reasoning_level_override: self.reasoning_level_override,
review_request,
role: self.role,
size: self.size,
speed_mode: self.speed_mode,
status: self.status,
summary,
title: self.title,
updated_at: self.updated_at,
}
}
fn into_session_list_row(
self,
review_request: Option<SessionReviewRequestRow>,
) -> SessionListRow {
SessionListRow {
added_lines: self.added_lines,
agent: self.agent,
base_branch: self.base_branch,
created_at: self.created_at,
deleted_lines: self.deleted_lines,
has_diff: self.has_diff,
id: self.id,
in_progress_started_at: self.in_progress_started_at,
in_progress_total_seconds: self.in_progress_total_seconds,
input_tokens: self.input_tokens,
is_draft: self.is_draft,
model: self.model,
output_tokens: self.output_tokens,
parent_session_id: self.parent_session_id,
personality_id: self.personality_id,
project_id: self.project_id,
published_upstream_ref: self.published_upstream_ref,
reasoning_level_override: self.reasoning_level_override,
review_request,
role: self.role,
size: self.size,
speed_mode: self.speed_mode,
status: self.status,
title: self.title,
updated_at: self.updated_at,
}
}
}
#[derive(sqlx::FromRow)]
struct SessionJoinRow {
added_lines: i64,
agent: String,
base_branch: String,
created_at: i64,
deleted_lines: i64,
has_diff: Option<bool>,
id: String,
in_progress_started_at: Option<i64>,
in_progress_total_seconds: i64,
input_tokens: i64,
is_draft: bool,
model: String,
output_tokens: i64,
parent_session_id: Option<String>,
personality_id: Option<String>,
project_id: Option<i64>,
prompt: String,
published_upstream_ref: Option<String>,
questions: Option<String>,
reasoning_level_override: Option<String>,
review_request_display_id: Option<String>,
review_request_forge_kind: Option<String>,
review_request_last_refreshed_at: Option<i64>,
review_request_source_branch: Option<String>,
review_request_state: Option<String>,
review_request_status_summary: Option<String>,
review_request_target_branch: Option<String>,
review_request_title: Option<String>,
review_request_web_url: Option<String>,
role: Option<String>,
size: String,
speed_mode: String,
status: String,
summary: Option<String>,
title: Option<String>,
updated_at: i64,
}
impl SessionJoinRow {
fn has_loadable_status(&self) -> bool {
if let Err(error) = status::validate_session(&self.status) {
warn!(
session_id = %self.id,
%error,
"Skipping session with invalid persisted status"
);
return false;
}
true
}
fn into_session_row(self) -> SessionRow {
let (metadata, detail, review_request) = self.into_parts();
metadata.into_session_row(
detail.prompt,
detail.questions,
detail.summary,
review_request,
)
}
fn into_session_list_row(self) -> SessionListRow {
let (metadata, _, review_request) = self.into_parts();
metadata.into_session_list_row(review_request)
}
fn into_parts(
self,
) -> (
SessionRowMetadata,
SessionDetailRow,
Option<SessionReviewRequestRow>,
) {
let Self {
added_lines,
agent,
base_branch,
created_at,
deleted_lines,
has_diff,
id,
in_progress_started_at,
in_progress_total_seconds,
input_tokens,
is_draft,
model,
output_tokens,
parent_session_id,
personality_id,
project_id,
prompt,
published_upstream_ref,
questions,
reasoning_level_override,
review_request_display_id,
review_request_forge_kind,
review_request_last_refreshed_at,
review_request_source_branch,
review_request_state,
review_request_status_summary,
review_request_target_branch,
review_request_title,
review_request_web_url,
role,
size,
speed_mode,
status,
summary,
title,
updated_at,
} = self;
let metadata = SessionRowMetadata {
added_lines,
agent,
base_branch,
created_at,
deleted_lines,
has_diff,
id,
in_progress_started_at,
in_progress_total_seconds,
input_tokens,
is_draft,
model,
output_tokens,
parent_session_id,
personality_id,
project_id,
published_upstream_ref,
reasoning_level_override,
role,
size,
speed_mode,
status,
title,
updated_at,
};
let detail = SessionDetailRow {
prompt,
questions,
summary,
};
let review_request = SessionReviewRequestJoinRow {
display_id: review_request_display_id,
forge_kind: review_request_forge_kind,
last_refreshed_at: review_request_last_refreshed_at,
source_branch: review_request_source_branch,
state: review_request_state,
status_summary: review_request_status_summary,
target_branch: review_request_target_branch,
title: review_request_title,
web_url: review_request_web_url,
}
.into_review_request_row();
(metadata, detail, review_request)
}
}
struct SessionReviewRequestJoinRow {
display_id: Option<String>,
forge_kind: Option<String>,
last_refreshed_at: Option<i64>,
source_branch: Option<String>,
state: Option<String>,
status_summary: Option<String>,
target_branch: Option<String>,
title: Option<String>,
web_url: Option<String>,
}
impl SessionReviewRequestJoinRow {
fn into_review_request_row(self) -> Option<SessionReviewRequestRow> {
let Self {
display_id,
forge_kind,
last_refreshed_at,
source_branch,
state,
status_summary,
target_branch,
title,
web_url,
} = self;
Some(SessionReviewRequestRow {
display_id: display_id?,
forge_kind: forge_kind?,
last_refreshed_at: last_refreshed_at?,
source_branch: source_branch?,
state: state?,
status_summary,
target_branch: target_branch?,
title: title?,
web_url: web_url?,
})
}
}
#[async_trait]
impl SessionRepository for SqliteSessionRepository {
async fn append_session_message(
&self,
id: &str,
kind: SessionMessageKind,
content: &str,
) -> Result<(), DbError> {
self.2.append(id, kind, content).await
}
async fn backfill_session_project(&self, project_id: i64) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET project_id = ?,
updated_at = ?
WHERE project_id IS NULL
",
project_id,
now
)
.execute(&self.0)
.await?;
Ok(())
}
async fn delete_session(&self, id: &str) -> Result<(), DbError> {
let now = self.now();
let mut transaction = self.0.begin().await?;
sqlx::query!(
r"
UPDATE session
SET parent_session_id = NULL,
base_branch = COALESCE((SELECT base_branch FROM session WHERE id = ?), base_branch),
updated_at = ?
WHERE parent_session_id = ?
AND status <> 'Canceled'
",
id,
now,
id
)
.execute(&mut *transaction)
.await?;
sqlx::query!(
r"
DELETE FROM session
WHERE id = ?
",
id
)
.execute(&mut *transaction)
.await?;
transaction.commit().await?;
Ok(())
}
async fn get_session_base_branch(&self, id: &str) -> Result<Option<String>, DbError> {
let row = sqlx::query_as!(
RequiredStringValueRow,
r#"
SELECT base_branch AS "value!: _"
FROM session
WHERE id = ?
"#,
id
)
.fetch_optional(&self.0)
.await?;
Ok(row.map(|row| row.value))
}
async fn get_session_parent_session_id(&self, id: &str) -> Result<Option<String>, DbError> {
let value = sqlx::query_scalar!(
r"
SELECT parent_session_id
FROM session
WHERE id = ?
",
id
)
.fetch_optional(&self.0)
.await?
.flatten();
Ok(value)
}
async fn get_session_stack_base_commit_hash(
&self,
id: &str,
) -> Result<Option<String>, DbError> {
let value = sqlx::query_scalar!(
r"
SELECT stack_base_commit_hash
FROM session
WHERE id = ?
",
id
)
.fetch_optional(&self.0)
.await?
.flatten();
Ok(value)
}
async fn get_session_instruction_conversation_id(
&self,
id: &str,
) -> Result<Option<String>, DbError> {
let row = sqlx::query_as!(
SessionInstructionStateRow,
r"
SELECT app_server_instruction_provider_conversation_id
FROM session
WHERE id = ?
",
id
)
.fetch_optional(&self.0)
.await?;
Ok(row.and_then(SessionInstructionStateRow::into_instruction_conversation_id))
}
async fn get_session_provider_conversation_id(
&self,
id: &str,
) -> Result<Option<String>, DbError> {
let value = sqlx::query_scalar!(
r"SELECT provider_conversation_id FROM session WHERE id = ?",
id
)
.fetch_optional(&self.0)
.await?
.flatten();
Ok(value)
}
async fn insert_draft_session(
&self,
id: &str,
model: &str,
base_branch: &str,
status: &str,
project_id: i64,
) -> Result<(), DbError> {
let agent = persisted_agent_for_model(model);
insert_session_with_draft_mode(
&self.0,
self.now(),
InsertSessionRow {
agent: &agent,
base_branch,
id,
is_draft: true,
model,
orchestration_task_id: None,
parent_session_id: None,
personality_id: None,
project_id,
reasoning_level: ReasoningLevel::default(),
role: None,
speed_mode: SpeedMode::Normal,
status,
},
)
.await
}
async fn insert_stacked_draft_session(
&self,
id: &str,
model: &str,
base_branch: &str,
status: &str,
parent_session_id: &str,
project_id: i64,
) -> Result<(), DbError> {
let agent = persisted_agent_for_model(model);
insert_session_with_draft_mode(
&self.0,
self.now(),
InsertSessionRow {
agent: &agent,
base_branch,
id,
is_draft: true,
model,
orchestration_task_id: None,
parent_session_id: Some(parent_session_id),
personality_id: None,
project_id,
reasoning_level: ReasoningLevel::default(),
role: None,
speed_mode: SpeedMode::Normal,
status,
},
)
.await
}
async fn insert_session(
&self,
id: &str,
model: &str,
base_branch: &str,
status: &str,
project_id: i64,
) -> Result<(), DbError> {
let agent = persisted_agent_for_model(model);
insert_session_with_draft_mode(
&self.0,
self.now(),
InsertSessionRow {
agent: &agent,
base_branch,
id,
is_draft: false,
model,
orchestration_task_id: None,
parent_session_id: None,
personality_id: None,
project_id,
reasoning_level: ReasoningLevel::default(),
role: None,
speed_mode: SpeedMode::Normal,
status,
},
)
.await
}
async fn insert_session_with_agent(
&self,
session: PersistedSessionCreation<'_>,
) -> Result<(), DbError> {
let PersistedSessionCreation {
agent,
base_branch,
id,
is_draft,
model,
orchestration_task_id,
parent_session_id,
personality_id,
project_id,
reasoning_level,
role,
speed_mode,
status,
} = session;
insert_session_with_draft_mode(
&self.0,
self.now(),
InsertSessionRow {
agent,
base_branch,
id,
is_draft,
model,
orchestration_task_id,
parent_session_id,
personality_id,
project_id,
reasoning_level,
role,
speed_mode,
status,
},
)
.await
}
async fn fork_session_snapshot(
&self,
snapshot: ForkSessionSnapshot<'_>,
) -> Result<(), DbError> {
self.3.fork(snapshot).await
}
async fn load_session(&self, session_id: &str) -> Result<Option<SessionRow>, DbError> {
let row = sqlx::query_as::<_, SessionJoinRow>(
r"
SELECT session.base_branch AS base_branch,
session.added_lines AS added_lines,
session.agent AS agent,
session.created_at AS created_at,
session.deleted_lines AS deleted_lines,
session.has_diff AS has_diff,
session.id AS id,
session.in_progress_started_at,
session.in_progress_total_seconds AS in_progress_total_seconds,
session.input_tokens AS input_tokens,
session.is_draft AS is_draft,
session.model AS model,
session.output_tokens AS output_tokens,
session.parent_session_id,
session.personality_id,
session.project_id,
session.prompt AS prompt,
session.reasoning_level AS reasoning_level_override,
session.speed_mode AS speed_mode,
session.published_upstream_ref,
session.questions,
session_review_request.display_id AS review_request_display_id,
session_review_request.forge_kind AS review_request_forge_kind,
session_review_request.last_refreshed_at AS review_request_last_refreshed_at,
session_review_request.source_branch AS review_request_source_branch,
session_review_request.state AS review_request_state,
session_review_request.status_summary AS review_request_status_summary,
session_review_request.target_branch AS review_request_target_branch,
session_review_request.title AS review_request_title,
session_review_request.web_url AS review_request_web_url,
session.role,
session.size AS size,
session.status AS status,
session.summary,
session.title,
session.updated_at AS updated_at
FROM session
LEFT JOIN session_review_request
ON session_review_request.session_id = session.id
WHERE session.id = ?
",
)
.bind(session_id)
.fetch_optional(&self.0)
.await?;
let row = row.map(SessionJoinRow::into_session_row);
if let Some(row) = &row {
status::validate_session(&row.status)?;
}
Ok(row)
}
async fn load_active_session_agent_models(&self) -> Result<Vec<SessionAgentModelRow>, DbError> {
let rows = sqlx::query_as::<_, SessionAgentModelRow>(
r"
SELECT agent,
id,
model,
status
FROM session
WHERE status NOT IN ('Merged', 'Done', 'Canceled')
ORDER BY id
",
)
.fetch_all(&self.0)
.await?;
Ok(rows)
}
#[cfg(any(test, feature = "test-utils"))]
async fn load_sessions(&self) -> Result<Vec<SessionRow>, DbError> {
let rows = sqlx::query_as!(
SessionJoinRow,
r#"
SELECT session.base_branch AS base_branch,
session.added_lines AS added_lines,
session.agent AS agent,
session.created_at AS created_at,
session.deleted_lines AS deleted_lines,
session.has_diff AS "has_diff: bool",
session.id AS id,
session.in_progress_started_at,
session.in_progress_total_seconds AS in_progress_total_seconds,
session.input_tokens AS input_tokens,
session.is_draft AS "is_draft: bool",
session.model AS model,
session.output_tokens AS output_tokens,
session.parent_session_id,
session.personality_id,
session.project_id,
session.prompt AS prompt,
session.reasoning_level AS reasoning_level_override,
session.speed_mode AS speed_mode,
session.published_upstream_ref,
session.questions,
session_review_request.display_id AS review_request_display_id,
session_review_request.forge_kind AS review_request_forge_kind,
session_review_request.last_refreshed_at AS review_request_last_refreshed_at,
session_review_request.source_branch AS review_request_source_branch,
session_review_request.state AS review_request_state,
session_review_request.status_summary AS review_request_status_summary,
session_review_request.target_branch AS review_request_target_branch,
session_review_request.title AS review_request_title,
session_review_request.web_url AS review_request_web_url,
session.role,
session.size AS size,
session.status AS status,
session.summary,
session.title,
session.updated_at AS updated_at
FROM session
LEFT JOIN session_review_request
ON session_review_request.session_id = session.id
ORDER BY session.updated_at DESC, session.created_at DESC, session.id
"#
)
.fetch_all(&self.0)
.await?;
let rows = rows
.into_iter()
.filter(SessionJoinRow::has_loadable_status)
.map(SessionJoinRow::into_session_row)
.collect::<Vec<_>>();
Ok(rows)
}
async fn load_sessions_for_project(
&self,
project_id: i64,
) -> Result<Vec<SessionListRow>, DbError> {
let rows = sqlx::query_as!(
SessionJoinRow,
r#"
SELECT session.base_branch AS base_branch,
session.added_lines AS added_lines,
session.agent AS agent,
session.created_at AS created_at,
session.deleted_lines AS deleted_lines,
session.has_diff AS "has_diff: bool",
session.id AS id,
session.in_progress_started_at,
session.in_progress_total_seconds AS in_progress_total_seconds,
session.input_tokens AS input_tokens,
session.is_draft AS "is_draft: bool",
session.model AS model,
session.output_tokens AS output_tokens,
session.parent_session_id,
session.personality_id,
session.project_id,
'' AS "prompt!: String",
session.reasoning_level AS reasoning_level_override,
session.speed_mode AS speed_mode,
session.published_upstream_ref,
NULL AS "questions: String",
session_review_request.display_id AS review_request_display_id,
session_review_request.forge_kind AS review_request_forge_kind,
session_review_request.last_refreshed_at AS review_request_last_refreshed_at,
session_review_request.source_branch AS review_request_source_branch,
session_review_request.state AS review_request_state,
session_review_request.status_summary AS review_request_status_summary,
session_review_request.target_branch AS review_request_target_branch,
session_review_request.title AS review_request_title,
session_review_request.web_url AS review_request_web_url,
session.role,
session.size AS size,
session.status AS status,
NULL AS "summary: String",
session.title,
session.updated_at AS updated_at
FROM session
LEFT JOIN session_review_request
ON session_review_request.session_id = session.id
WHERE session.project_id = ?
ORDER BY session.updated_at DESC, session.created_at DESC, session.id
"#,
project_id
)
.fetch_all(&self.0)
.await?;
let rows = rows
.into_iter()
.filter(SessionJoinRow::has_loadable_status)
.map(SessionJoinRow::into_session_list_row)
.collect::<Vec<_>>();
Ok(rows)
}
async fn load_session_detail(
&self,
session_id: &str,
) -> Result<Option<SessionDetailRow>, DbError> {
let row = sqlx::query_as!(
SessionDetailRow,
r"
SELECT prompt,
questions,
summary
FROM session
WHERE id = ?
",
session_id
)
.fetch_optional(&self.0)
.await?;
Ok(row)
}
async fn load_session_messages(
&self,
session_id: &str,
) -> Result<Vec<SessionMessageRow>, DbError> {
let rows = sqlx::query_as!(
SessionMessageRow,
r"
SELECT content,
kind,
position
FROM session_message
WHERE session_id = ?
ORDER BY position, id
",
session_id
)
.fetch_all(&self.0)
.await?;
Ok(rows)
}
async fn load_session_focused_reviews_for_project(
&self,
project_id: i64,
) -> Result<Vec<SessionFocusedReviewRow>, DbError> {
let rows = sqlx::query_as!(
SessionFocusedReviewRow,
r#"
SELECT id AS session_id,
focused_review_diff_hash AS "diff_hash!: String",
focused_review_text AS "text!: String"
FROM session
WHERE project_id = ?
AND focused_review_diff_hash IS NOT NULL
AND focused_review_text IS NOT NULL
AND focused_review_text <> ''
ORDER BY updated_at DESC, id
"#,
project_id
)
.fetch_all(&self.0)
.await?;
Ok(rows)
}
async fn load_sessions_metadata(&self) -> Result<(i64, i64), DbError> {
let row = sqlx::query_as!(
SessionStatsMetadataRow,
r#"
SELECT (SELECT COUNT(*) FROM session) AS "session_count!: _",
COALESCE(
(
SELECT updated_at
FROM session
ORDER BY updated_at DESC, id
LIMIT 1
),
0
) AS "max_updated_at!: _"
"#
)
.fetch_one(&self.0)
.await?;
Ok((row.session_count, row.max_updated_at))
}
async fn load_session_project_id(&self, session_id: &str) -> Result<Option<i64>, DbError> {
let row = sqlx::query_as!(
OptionalI64ValueRow,
r#"
SELECT project_id AS "value: _"
FROM session
WHERE id = ?
"#,
session_id
)
.fetch_optional(&self.0)
.await?;
Ok(row.and_then(|row| row.value))
}
async fn load_session_personality_state(
&self,
session_id: &str,
) -> Result<Option<SessionPersonalityState>, DbError> {
let state = sqlx::query_as!(
SessionPersonalityState,
r"
SELECT applied_personality_id,
applied_personality_prompt_hash,
personality_id
FROM session
WHERE id = ?
",
session_id
)
.fetch_optional(&self.0)
.await?;
Ok(state)
}
async fn load_pending_stack_restack_session_ids(
&self,
project_id: i64,
) -> Result<Vec<String>, DbError> {
let session_ids = sqlx::query_scalar!(
r"
SELECT id
FROM session
WHERE project_id = ?
AND parent_session_id IS NULL
AND stack_base_commit_hash IS NOT NULL
AND status IN ('Review', 'AgentReview')
ORDER BY updated_at ASC, id ASC
",
project_id
)
.fetch_all(&self.0)
.await?;
Ok(session_ids)
}
async fn load_session_published_upstream_ref(
&self,
id: &str,
) -> Result<Option<String>, DbError> {
let value = sqlx::query_scalar!(
r"SELECT published_upstream_ref FROM session WHERE id = ?",
id
)
.fetch_optional(&self.0)
.await?
.flatten();
Ok(value)
}
async fn load_session_merged_commit_hash(
&self,
session_id: &str,
) -> Result<Option<String>, DbError> {
let row = sqlx::query_scalar!(
r"
SELECT merged_commit_hash
FROM session
WHERE id = ?
",
session_id
)
.fetch_optional(&self.0)
.await?;
Ok(row.flatten())
}
async fn load_session_archived_diff(
&self,
session_id: &str,
) -> Result<Option<String>, DbError> {
let row = sqlx::query_scalar!(
r"
SELECT archived_diff
FROM session
WHERE id = ?
",
session_id
)
.fetch_optional(&self.0)
.await?;
Ok(row.flatten())
}
async fn load_session_reasoning_level(
&self,
session_id: &str,
) -> Result<ReasoningLevel, DbError> {
let value = sqlx::query_scalar!(
r"SELECT reasoning_level FROM session WHERE id = ?",
session_id
)
.fetch_optional(&self.0)
.await?
.flatten();
Ok(value
.and_then(|value| value.parse::<ReasoningLevel>().ok())
.unwrap_or_default())
}
async fn load_session_speed_mode(&self, session_id: &str) -> Result<SpeedMode, DbError> {
let value = sqlx::query_scalar!(r"SELECT speed_mode FROM session WHERE id = ?", session_id)
.fetch_optional(&self.0)
.await?;
Ok(value
.and_then(|value| value.parse::<SpeedMode>().ok())
.unwrap_or_default())
}
async fn restack_child_sessions_after_parent_merge(
&self,
parent_session_id: &str,
base_branch: &str,
parent_commit_hash: Option<String>,
) -> Result<Vec<String>, DbError> {
let now = self.now();
let mut transaction = self.0.begin().await?;
let materialized_child_ids = sqlx::query_scalar!(
r"
SELECT id
FROM session
WHERE parent_session_id = ?
AND status NOT IN ('Canceled', 'Draft')
ORDER BY created_at ASC, id ASC
",
parent_session_id
)
.fetch_all(&mut *transaction)
.await?;
sqlx::query!(
r"
UPDATE session
SET parent_session_id = NULL,
base_branch = ?,
stack_base_commit_hash = CASE
WHEN status = 'Draft' THEN NULL
ELSE COALESCE(stack_base_commit_hash, ?)
END,
updated_at = ?
WHERE parent_session_id = ?
AND status <> 'Canceled'
",
base_branch,
parent_commit_hash,
now,
parent_session_id
)
.execute(&mut *transaction)
.await?;
transaction.commit().await?;
Ok(materialized_child_ids)
}
async fn load_session_summary(&self, session_id: &str) -> Result<Option<String>, DbError> {
let row = sqlx::query_scalar!(
r"
SELECT summary
FROM session
WHERE id = ?
",
session_id
)
.fetch_optional(&self.0)
.await?;
Ok(row.flatten())
}
async fn load_session_timestamps(
&self,
session_id: &str,
) -> Result<Option<(i64, i64)>, DbError> {
let row = sqlx::query_as!(
SessionTimestampsRow,
r#"
SELECT created_at, updated_at
FROM session
WHERE id = ?
"#,
session_id
)
.fetch_optional(&self.0)
.await?;
Ok(row.map(|row| (row.created_at, row.updated_at)))
}
async fn persist_session_turn_metadata(
&self,
session_id: &str,
turn_metadata: &SessionTurnMetadata,
) -> Result<(), DbError> {
let now = self.now();
let mut transaction = self.0.begin().await?;
let session_update = sqlx::query!(
r"
UPDATE session
SET questions = ?,
summary = ?,
provider_conversation_id = ?,
app_server_instruction_provider_conversation_id = ?,
applied_personality_id = ?,
applied_personality_prompt_hash = ?,
updated_at = ?
WHERE id = ?
",
turn_metadata.questions_json.as_str(),
turn_metadata.summary.as_str(),
turn_metadata.provider_conversation_id.as_deref(),
turn_metadata.instruction_conversation_id.as_deref(),
turn_metadata.applied_personality_id.as_deref(),
turn_metadata.applied_personality_prompt_hash.as_deref(),
now,
session_id
)
.execute(&mut *transaction)
.await?;
if session_update.rows_affected() != 1 {
return Err(sqlx::Error::RowNotFound.into());
}
if turn_metadata.token_usage_delta.input_tokens != 0
|| turn_metadata.token_usage_delta.output_tokens != 0
{
sqlx::query!(
r"
UPDATE session
SET input_tokens = input_tokens + ?,
output_tokens = output_tokens + ?,
updated_at = ?
WHERE id = ?
",
turn_metadata.token_usage_delta.input_tokens.cast_signed(),
turn_metadata.token_usage_delta.output_tokens.cast_signed(),
now,
session_id
)
.execute(&mut *transaction)
.await?;
sqlx::query!(
r"
INSERT INTO session_usage (
session_id, model, created_at, input_tokens, output_tokens, invocation_count
)
VALUES (?, ?, ?, ?, ?, 1)
ON CONFLICT(session_id, model) DO UPDATE SET
input_tokens = input_tokens + excluded.input_tokens,
output_tokens = output_tokens + excluded.output_tokens,
invocation_count = invocation_count + 1
",
session_id,
turn_metadata.model.as_str(),
now,
turn_metadata.token_usage_delta.input_tokens.cast_signed(),
turn_metadata.token_usage_delta.output_tokens.cast_signed()
)
.execute(&mut *transaction)
.await?;
}
transaction.commit().await?;
Ok(())
}
async fn update_session_diff_stats(
&self,
added_lines: u64,
deleted_lines: u64,
has_diff: bool,
id: &str,
size: &str,
) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET added_lines = ?,
deleted_lines = ?,
has_diff = ?,
size = ?,
updated_at = ?
WHERE id = ?
AND (
added_lines <> ?
OR deleted_lines <> ?
OR has_diff IS NOT ?
OR size <> ?
)
",
added_lines.cast_signed(),
deleted_lines.cast_signed(),
has_diff,
size,
now,
id,
added_lines.cast_signed(),
deleted_lines.cast_signed(),
has_diff,
size
)
.execute(&self.0)
.await?;
Ok(())
}
async fn mark_session_diff_unknown(&self, id: &str) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET has_diff = NULL,
updated_at = ?
WHERE id = ?
AND has_diff IS NOT NULL
",
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_instruction_conversation_id(
&self,
id: &str,
provider_conversation_id: Option<String>,
) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET app_server_instruction_provider_conversation_id = ?,
updated_at = ?
WHERE id = ?
",
provider_conversation_id.as_deref(),
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_model(&self, id: &str, model: &str) -> Result<(), DbError> {
let agent = persisted_agent_for_model(model);
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET agent = ?,
model = ?,
updated_at = ?
WHERE id = ?
",
agent,
model,
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_personality_id(
&self,
id: &str,
personality_id: Option<String>,
) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET personality_id = ?,
updated_at = ?
WHERE id = ?
",
personality_id.as_deref(),
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_agent_model(
&self,
id: &str,
agent: &str,
model: &str,
) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET agent = ?,
model = ?,
updated_at = ?
WHERE id = ?
",
agent,
model,
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_active_session_agent_model(
&self,
id: &str,
agent: &str,
model: &str,
) -> Result<(), DbError> {
sqlx::query!(
r"
UPDATE session
SET agent = ?,
model = ?
WHERE id = ?
AND status NOT IN ('Merged', 'Done', 'Canceled')
",
agent,
model,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn clear_session_draft_flag(&self, id: &str) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET is_draft = 0,
updated_at = ?
WHERE id = ?
",
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_merged_commit_hash(
&self,
id: &str,
merged_commit_hash: Option<String>,
) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET merged_commit_hash = ?,
updated_at = ?
WHERE id = ?
",
merged_commit_hash.as_deref(),
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_archived_diff(
&self,
id: &str,
archived_diff: Option<String>,
) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET archived_diff = ?,
updated_at = ?
WHERE id = ?
",
archived_diff.as_deref(),
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_stack_base_commit_hash(
&self,
id: &str,
stack_base_commit_hash: Option<String>,
) -> Result<(), DbError> {
let now = self.now();
sqlx::query(
r"
UPDATE session
SET stack_base_commit_hash = ?,
updated_at = ?
WHERE id = ?
",
)
.bind(stack_base_commit_hash)
.bind(now)
.bind(id)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_prompt(&self, id: &str, prompt: &str) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET prompt = ?,
updated_at = ?
WHERE id = ?
",
prompt,
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_provider_conversation_id(
&self,
id: &str,
provider_conversation_id: Option<String>,
) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET provider_conversation_id = ?,
updated_at = ?
WHERE id = ?
",
provider_conversation_id.as_deref(),
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_questions(&self, id: &str, questions: &str) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET questions = ?,
updated_at = ?
WHERE id = ?
",
questions,
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_reasoning_level(
&self,
id: &str,
reasoning_level: ReasoningLevel,
) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r#"
UPDATE session
SET reasoning_level = ?,
updated_at = ?
WHERE id = ?
"#,
reasoning_level.as_str(),
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_speed_mode(
&self,
id: &str,
speed_mode: SpeedMode,
) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r#"
UPDATE session
SET speed_mode = ?,
updated_at = ?
WHERE id = ?
"#,
speed_mode.as_str(),
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_published_upstream_ref(
&self,
id: &str,
published_upstream_ref: Option<String>,
) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET published_upstream_ref = ?,
updated_at = ?
WHERE id = ?
",
published_upstream_ref.as_deref(),
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_stats(&self, id: &str, stats: &SessionStats) -> Result<(), DbError> {
if stats.input_tokens == 0 && stats.output_tokens == 0 {
return Ok(());
}
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET input_tokens = input_tokens + ?,
output_tokens = output_tokens + ?,
updated_at = ?
WHERE id = ?
",
stats.input_tokens.cast_signed(),
stats.output_tokens.cast_signed(),
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_status_with_timing_at(
&self,
id: &str,
status: &str,
timestamp_seconds: i64,
) -> Result<(), DbError> {
status::validate_session(status)?;
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET status = ?,
in_progress_total_seconds = CASE
WHEN ? = 'InProgress' OR in_progress_started_at IS NULL THEN in_progress_total_seconds
ELSE in_progress_total_seconds + MAX(0, ? - in_progress_started_at)
END,
in_progress_started_at = CASE
WHEN ? = 'InProgress' THEN COALESCE(in_progress_started_at, ?)
ELSE NULL
END,
updated_at = ?
WHERE id = ?
",
status,
status,
timestamp_seconds,
status,
timestamp_seconds,
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_summary(&self, id: &str, summary: &str) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET summary = ?,
updated_at = ?
WHERE id = ?
",
summary,
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_focused_review(
&self,
id: &str,
status: Option<FocusedReviewStatus>,
diff_hash: Option<String>,
text: Option<String>,
) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r"
UPDATE session
SET focused_review_status = ?,
focused_review_diff_hash = ?,
focused_review_text = ?,
updated_at = ?
WHERE id = ?
",
status.map(|status| status.to_string()),
diff_hash.as_deref(),
text.as_deref(),
now,
id
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_title(&self, id: &str, title: &str) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r#"
UPDATE session
SET title = ?,
is_title_provisional = 0,
title_generation = title_generation + 1,
applied_title_generation = title_generation + 1,
updated_at = ?
WHERE id = ?
"#,
title,
now,
id,
)
.execute(&self.0)
.await?;
Ok(())
}
async fn update_session_provisional_title(&self, id: &str, title: &str) -> Result<(), DbError> {
let now = self.now();
sqlx::query!(
r#"
UPDATE session
SET title = ?,
is_title_provisional = 1,
title_generation = title_generation + 1,
applied_title_generation = title_generation + 1,
updated_at = ?
WHERE id = ?
"#,
title,
now,
id,
)
.execute(&self.0)
.await?;
Ok(())
}
async fn begin_session_title_generation(
&self,
id: &str,
requires_provisional_title: bool,
) -> Result<Option<i64>, DbError> {
let now = self.now();
let generation = if requires_provisional_title {
sqlx::query_scalar!(
r#"
UPDATE session
SET is_title_provisional = 1,
title_generation = title_generation + 1,
updated_at = ?
WHERE id = ?
AND is_title_provisional = 1
RETURNING title_generation
"#,
now,
id,
)
.fetch_optional(&self.0)
.await?
} else {
sqlx::query_scalar!(
r#"
UPDATE session
SET is_title_provisional = 1,
title_generation = title_generation + 1,
updated_at = ?
WHERE id = ?
RETURNING title_generation
"#,
now,
id,
)
.fetch_optional(&self.0)
.await?
};
Ok(generation)
}
async fn update_session_title_for_generation(
&self,
id: &str,
expected_generation: i64,
title: &str,
) -> Result<bool, DbError> {
let now = self.now();
let result = sqlx::query!(
r#"
UPDATE session
SET title = ?,
is_title_provisional = 0,
applied_title_generation = ?,
updated_at = ?
WHERE id = ?
AND title_generation >= ?
AND applied_title_generation < ?
"#,
title,
expected_generation,
now,
id,
expected_generation,
expected_generation,
)
.execute(&self.0)
.await?;
Ok(result.rows_affected() > 0)
}
#[cfg(any(test, feature = "test-utils"))]
async fn update_session_created_at(&self, id: &str, created_at: i64) -> Result<(), DbError> {
sqlx::query!(
r"
UPDATE session
SET created_at = ?
WHERE id = ?
",
created_at,
id
)
.execute(&self.0)
.await?;
Ok(())
}
#[cfg(any(test, feature = "test-utils"))]
async fn update_session_updated_at(&self, id: &str, updated_at: i64) -> Result<(), DbError> {
sqlx::query!(
r"
UPDATE session
SET updated_at = ?
WHERE id = ?
",
updated_at,
id
)
.execute(&self.0)
.await?;
Ok(())
}
}
struct InsertSessionRow<'a> {
agent: &'a str,
base_branch: &'a str,
id: &'a str,
is_draft: bool,
model: &'a str,
orchestration_task_id: Option<i64>,
parent_session_id: Option<&'a str>,
personality_id: Option<&'a str>,
project_id: i64,
reasoning_level: ReasoningLevel,
role: Option<&'a str>,
speed_mode: SpeedMode,
status: &'a str,
}
async fn insert_session_with_draft_mode(
pool: &SqlitePool,
timestamp_seconds: i64,
row: InsertSessionRow<'_>,
) -> Result<(), DbError> {
let InsertSessionRow {
agent,
base_branch,
id,
is_draft,
model,
orchestration_task_id,
parent_session_id,
personality_id,
project_id,
reasoning_level,
role,
speed_mode,
status,
} = row;
status::validate_session(status)?;
sqlx::query(
r"
INSERT INTO session (
id,
agent,
model,
base_branch,
status,
has_diff,
is_draft,
parent_session_id,
personality_id,
project_id,
reasoning_level,
role,
speed_mode,
orchestration_task_id,
prompt,
created_at,
updated_at
)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
",
)
.bind(id)
.bind(agent)
.bind(model)
.bind(base_branch)
.bind(status)
.bind(Option::<bool>::None)
.bind(is_draft)
.bind(parent_session_id)
.bind(personality_id)
.bind(project_id)
.bind(reasoning_level.as_str())
.bind(role)
.bind(speed_mode.as_str())
.bind(orchestration_task_id)
.bind("")
.bind(timestamp_seconds)
.bind(timestamp_seconds)
.execute(pool)
.await?;
Ok(())
}
fn persisted_agent_for_model(model: &str) -> String {
AgentModel::parse_persisted(model).map_or_else(
|_| persisted_agent_for_unknown_model(model).to_string(),
|agent_model| persisted_agent_for_known_model(model, agent_model).to_string(),
)
}
fn persisted_agent_for_known_model(model: &str, agent_model: AgentModel) -> AgentKind {
if model.starts_with("claude-") {
return AgentKind::Claude;
}
if model.starts_with("gpt-") {
return AgentKind::Codex;
}
if model.starts_with("gemini-") {
return AgentKind::Antigravity;
}
AgentKind::ALL
.iter()
.copied()
.find(|agent_kind| agent_kind.supports_model(agent_model))
.unwrap_or(AgentKind::Antigravity)
}
fn persisted_agent_for_unknown_model(model: &str) -> AgentKind {
if model.starts_with("claude-") {
return AgentKind::Claude;
}
if model.starts_with("gpt-") {
return AgentKind::Codex;
}
if model.starts_with("gemini-") {
return AgentKind::Antigravity;
}
AgentKind::Antigravity
}
#[cfg(test)]
mod tests {
use ag_session::{ForgeKind, ReviewRequest, ReviewRequestState, ReviewRequestSummary};
use super::*;
use crate::AppRepositories;
struct ForkResetRow {
applied_personality_id: Option<String>,
applied_personality_prompt_hash: Option<String>,
app_server_instruction_provider_conversation_id: Option<String>,
focused_review_diff_hash: Option<String>,
focused_review_text: Option<String>,
in_progress_started_at: Option<i64>,
in_progress_total_seconds: i64,
is_draft: bool,
merged_commit_hash: Option<String>,
parent_session_id: Option<String>,
provider_conversation_id: Option<String>,
published_upstream_ref: Option<String>,
questions: Option<String>,
stack_base_commit_hash: Option<String>,
}
impl SessionJoinRow {
fn fixture_for_test() -> Self {
Self {
added_lines: 14,
agent: "codex".to_string(),
base_branch: "main".to_string(),
created_at: 100,
deleted_lines: 6,
has_diff: Some(true),
id: "session-a".to_string(),
in_progress_started_at: None,
in_progress_total_seconds: 0,
input_tokens: 11,
is_draft: false,
model: "gpt-5.6-sol".to_string(),
output_tokens: 29,
parent_session_id: Some("parent-session".to_string()),
personality_id: Some("reviewer".to_string()),
project_id: Some(7),
prompt: "Implement feature".to_string(),
published_upstream_ref: Some("origin/session-a".to_string()),
questions: Some("Question text".to_string()),
reasoning_level_override: None,
review_request_display_id: Some("#42".to_string()),
review_request_forge_kind: Some("GitHub".to_string()),
review_request_last_refreshed_at: Some(456),
review_request_source_branch: Some("feature/forge".to_string()),
review_request_state: Some("Open".to_string()),
review_request_status_summary: Some("2 approvals, checks passing".to_string()),
review_request_target_branch: Some("main".to_string()),
review_request_title: Some("Add forge review support".to_string()),
review_request_web_url: Some(
"https://github.com/agentty-xyz/agentty/pull/42".to_string(),
),
role: Some("Orchestrator".to_string()),
size: "M".to_string(),
speed_mode: "normal".to_string(),
status: "Review".to_string(),
summary: Some("Summary text".to_string()),
title: Some("Review session".to_string()),
updated_at: 200,
}
}
}
fn expected_review_request_row() -> SessionReviewRequestRow {
SessionReviewRequestRow {
display_id: "#42".to_string(),
forge_kind: "GitHub".to_string(),
last_refreshed_at: 456,
source_branch: "feature/forge".to_string(),
state: "Open".to_string(),
status_summary: Some("2 approvals, checks passing".to_string()),
target_branch: "main".to_string(),
title: "Add forge review support".to_string(),
web_url: "https://github.com/agentty-xyz/agentty/pull/42".to_string(),
}
}
fn review_request_fixture() -> ReviewRequest {
ReviewRequest {
last_refreshed_at: 456,
summary: ReviewRequestSummary {
display_id: "#42".to_string(),
forge_kind: ForgeKind::GitHub,
source_branch: "feature/forge".to_string(),
state: ReviewRequestState::Open,
status_summary: Some("2 approvals, checks passing".to_string()),
target_branch: "main".to_string(),
title: "Add forge review support".to_string(),
web_url: "https://github.com/agentty-xyz/agentty/pull/42".to_string(),
},
}
}
async fn load_fork_reset_row(pool: &SqlitePool, session_id: &str) -> ForkResetRow {
sqlx::query_as!(
ForkResetRow,
r#"
SELECT app_server_instruction_provider_conversation_id,
applied_personality_id,
applied_personality_prompt_hash,
focused_review_diff_hash,
focused_review_text,
in_progress_started_at,
in_progress_total_seconds,
is_draft AS "is_draft: bool",
merged_commit_hash,
parent_session_id,
provider_conversation_id,
published_upstream_ref,
questions,
stack_base_commit_hash
FROM session
WHERE id = ?
"#,
session_id
)
.fetch_one(pool)
.await
.expect("failed to load fork reset row")
}
async fn seed_fork_snapshot_source(
database: &AppRepositories,
pool: &SqlitePool,
) -> (ForkResetRow, Option<SessionReviewRequestRow>) {
let project_id = database
.projects()
.upsert_project("/tmp/project", None)
.await
.expect("failed to upsert project");
database
.sessions()
.insert_session(
"parent-session",
"gpt-5.6-sol",
"main",
"Review",
project_id,
)
.await
.expect("failed to insert parent session");
database
.sessions()
.insert_stacked_draft_session(
"source-session",
"gpt-5.6-sol",
"wt/parent",
"Review",
"parent-session",
project_id,
)
.await
.expect("failed to insert source session");
seed_fork_snapshot_source_linkage(database).await;
seed_fork_snapshot_source_timing(database, pool).await;
let source_reset_row = load_fork_reset_row(pool, "source-session").await;
let source_review_request = database
.reviews()
.load_session_review_request("source-session")
.await
.expect("failed to load source review request");
(source_reset_row, source_review_request)
}
async fn seed_fork_snapshot_source_linkage(database: &AppRepositories) {
database
.sessions()
.update_session_personality_id("source-session", Some("reviewer".to_string()))
.await
.expect("failed to update personality id");
database
.sessions()
.persist_session_turn_metadata(
"source-session",
&SessionTurnMetadata {
applied_personality_id: Some("reviewer".to_string()),
applied_personality_prompt_hash: Some("personality-hash".to_string()),
instruction_conversation_id: None,
model: "gpt-5.6-sol".to_string(),
provider_conversation_id: None,
questions_json: "[]".to_string(),
summary: String::new(),
token_usage_delta: SessionStats::default(),
},
)
.await
.expect("failed to persist applied personality");
database
.sessions()
.update_session_provider_conversation_id(
"source-session",
Some("provider-thread".to_string()),
)
.await
.expect("failed to update provider conversation id");
database
.sessions()
.update_session_instruction_conversation_id(
"source-session",
Some("instruction-thread".to_string()),
)
.await
.expect("failed to update instruction conversation id");
database
.sessions()
.update_session_questions("source-session", r#"["Need detail?"]"#)
.await
.expect("failed to update questions");
database
.sessions()
.update_session_published_upstream_ref(
"source-session",
Some("origin/wt/source-session".to_string()),
)
.await
.expect("failed to update published upstream ref");
database
.sessions()
.update_session_merged_commit_hash("source-session", Some("merged123".to_string()))
.await
.expect("failed to update merged commit hash");
database
.sessions()
.update_session_focused_review(
"source-session",
Some(FocusedReviewStatus::Ready),
Some("diff123".to_string()),
Some("Focused review text".to_string()),
)
.await
.expect("failed to update focused review");
database
.sessions()
.update_session_stack_base_commit_hash(
"source-session",
Some("stackbase123".to_string()),
)
.await
.expect("failed to update stack base commit hash");
database
.sessions()
.update_session_stats(
"source-session",
&SessionStats {
added_lines: 0,
deleted_lines: 0,
diff_state: agent::SessionDiffState::Unknown,
input_tokens: 11,
output_tokens: 29,
},
)
.await
.expect("failed to update token stats");
database
.sessions()
.update_session_diff_stats(7, 3, true, "source-session", "S")
.await
.expect("failed to update source diff stats");
database
.reviews()
.update_session_review_request("source-session", Some(review_request_fixture()))
.await
.expect("failed to update review request");
}
async fn seed_fork_snapshot_source_timing(database: &AppRepositories, pool: &SqlitePool) {
database
.sessions()
.update_session_status_with_timing_at("source-session", "InProgress", 100)
.await
.expect("failed to open timing interval");
sqlx::query!(
r"
UPDATE session
SET in_progress_total_seconds = ?
WHERE id = ?
",
75_i64,
"source-session"
)
.execute(pool)
.await
.expect("failed to seed elapsed timing");
}
fn assert_source_reset_state(
source_row: &SessionRow,
source_reset_row: &ForkResetRow,
source_review_request: Option<&SessionReviewRequestRow>,
) {
assert_eq!(source_row.added_lines, 7);
assert_eq!(source_row.deleted_lines, 3);
assert_eq!(source_row.has_diff, Some(true));
assert_eq!(source_row.size, "S");
assert!(source_reset_row.is_draft);
assert_eq!(source_row.personality_id.as_deref(), Some("reviewer"));
assert_eq!(
source_reset_row.applied_personality_id.as_deref(),
Some("reviewer")
);
assert_eq!(
source_reset_row.applied_personality_prompt_hash.as_deref(),
Some("personality-hash")
);
assert_eq!(
source_reset_row.parent_session_id.as_deref(),
Some("parent-session")
);
assert_eq!(
source_reset_row.provider_conversation_id.as_deref(),
Some("provider-thread")
);
assert_eq!(
source_reset_row
.app_server_instruction_provider_conversation_id
.as_deref(),
Some("instruction-thread")
);
assert_eq!(
source_reset_row.published_upstream_ref.as_deref(),
Some("origin/wt/source-session")
);
assert_eq!(
source_reset_row.questions.as_deref(),
Some(r#"["Need detail?"]"#)
);
assert_eq!(
source_reset_row.merged_commit_hash.as_deref(),
Some("merged123")
);
assert_eq!(
source_reset_row.focused_review_diff_hash.as_deref(),
Some("diff123")
);
assert_eq!(
source_reset_row.focused_review_text.as_deref(),
Some("Focused review text")
);
assert_eq!(
source_reset_row.stack_base_commit_hash.as_deref(),
Some("stackbase123")
);
assert_eq!(source_reset_row.in_progress_started_at, Some(100));
assert_eq!(source_reset_row.in_progress_total_seconds, 75);
assert_eq!(
source_review_request.map(|review_request| review_request.display_id.as_str()),
Some("#42")
);
}
fn assert_fork_reset_state(
fork_row: &SessionRow,
fork_reset_row: &ForkResetRow,
fork_review_request: Option<&SessionReviewRequestRow>,
) {
assert_eq!(fork_row.status, "Review");
assert!(!fork_row.is_draft);
assert_eq!(fork_row.parent_session_id, None);
assert_eq!(fork_row.personality_id.as_deref(), Some("reviewer"));
assert_eq!(fork_row.input_tokens, 0);
assert_eq!(fork_row.output_tokens, 0);
assert_eq!(fork_row.added_lines, 0);
assert_eq!(fork_row.deleted_lines, 0);
assert_eq!(fork_row.has_diff, None);
assert_eq!(fork_row.size, "XS");
assert_eq!(fork_row.questions, None);
assert_eq!(fork_row.published_upstream_ref, None);
assert_eq!(fork_row.review_request, None);
assert_eq!(fork_reset_row.provider_conversation_id, None);
assert_eq!(fork_reset_row.applied_personality_id, None);
assert_eq!(fork_reset_row.applied_personality_prompt_hash, None);
assert_eq!(
fork_reset_row.app_server_instruction_provider_conversation_id,
None
);
assert_eq!(fork_reset_row.merged_commit_hash, None);
assert_eq!(fork_reset_row.focused_review_diff_hash, None);
assert_eq!(fork_reset_row.focused_review_text, None);
assert_eq!(fork_reset_row.questions, None);
assert_eq!(fork_reset_row.stack_base_commit_hash, None);
assert_eq!(fork_reset_row.in_progress_started_at, None);
assert_eq!(fork_reset_row.in_progress_total_seconds, 0);
assert_eq!(fork_review_request, None);
}
#[tokio::test]
async fn test_load_session_rejects_unknown_status() {
let (database, pool) = AppRepositories::in_memory_with_pool()
.await
.expect("db should open");
let project_id = database
.projects()
.upsert_project("/tmp/invalid-session", None)
.await
.expect("failed to upsert project");
database
.sessions()
.insert_session("session-a", "gpt-5.6-sol", "main", "Draft", project_id)
.await
.expect("failed to insert session");
sqlx::query("UPDATE session SET status = 'Unknown' WHERE id = 'session-a'")
.execute(&pool)
.await
.expect("failed to corrupt session status");
let result = database.sessions().load_session("session-a").await;
assert!(matches!(
result,
Err(DbError::InvalidStatus {
entity: "session",
value,
}) if value == "Unknown"
));
}
#[tokio::test]
async fn test_load_session_collections_skip_unknown_status() {
let (database, pool) = AppRepositories::in_memory_with_pool()
.await
.expect("db should open");
let project_id = database
.projects()
.upsert_project("/tmp/invalid-session-list", None)
.await
.expect("failed to upsert project");
for session_id in ["session-valid", "session-invalid"] {
database
.sessions()
.insert_session(session_id, "gpt-5.6-sol", "main", "Draft", project_id)
.await
.expect("failed to insert session");
}
sqlx::query("UPDATE session SET status = 'Unknown' WHERE id = 'session-invalid'")
.execute(&pool)
.await
.expect("failed to corrupt session status");
let all_sessions = database
.sessions()
.load_sessions()
.await
.expect("failed to load all sessions");
let project_sessions = database
.sessions()
.load_sessions_for_project(project_id)
.await
.expect("failed to load project sessions");
assert_eq!(
all_sessions
.iter()
.map(|session| session.id.as_str())
.collect::<Vec<_>>(),
["session-valid"]
);
assert_eq!(
project_sessions
.iter()
.map(|session| session.id.as_str())
.collect::<Vec<_>>(),
["session-valid"]
);
}
#[tokio::test]
async fn test_insert_session_starts_with_unknown_diff() {
let (database, _) = AppRepositories::in_memory_with_pool()
.await
.expect("db should open");
let project_id = database
.projects()
.upsert_project("/tmp/project", None)
.await
.expect("failed to upsert project");
database
.sessions()
.insert_session("session-a", "gpt-5.6-sol", "main", "Draft", project_id)
.await
.expect("failed to insert session");
let session = database
.sessions()
.load_sessions()
.await
.expect("failed to load sessions")
.into_iter()
.next()
.expect("missing inserted session");
assert_eq!(session.has_diff, None);
}
#[tokio::test]
async fn test_load_sessions_uses_created_at_to_break_updated_at_ties() {
let (database, pool) = AppRepositories::in_memory_with_pool()
.await
.expect("db should open");
let project_id = database
.projects()
.upsert_project("/tmp/project", None)
.await
.expect("failed to upsert project");
for session_id in ["a-older", "z-newer"] {
database
.sessions()
.insert_session(session_id, "gpt-5.6-sol", "main", "Review", project_id)
.await
.expect("failed to insert session");
}
sqlx::query!(
r"
UPDATE session
SET created_at = CASE id WHEN 'a-older' THEN 100 ELSE 200 END,
updated_at = 300
WHERE id IN ('a-older', 'z-newer')
"
)
.execute(&pool)
.await
.expect("failed to set session timestamps");
let all_session_ids = database
.sessions()
.load_sessions()
.await
.expect("failed to load sessions")
.into_iter()
.map(|session| session.id)
.collect::<Vec<_>>();
let project_session_ids = database
.sessions()
.load_sessions_for_project(project_id)
.await
.expect("failed to load project sessions")
.into_iter()
.map(|session| session.id)
.collect::<Vec<_>>();
assert_eq!(all_session_ids, ["z-newer", "a-older"]);
assert_eq!(project_session_ids, ["z-newer", "a-older"]);
}
#[tokio::test]
async fn test_fork_session_snapshot_resets_source_specific_state() {
let (database, pool) = AppRepositories::in_memory_with_pool()
.await
.expect("db should open");
let (source_reset_row, source_review_request) =
seed_fork_snapshot_source(&database, &pool).await;
database
.sessions()
.fork_session_snapshot(ForkSessionSnapshot {
new_session_id: "fork-session",
source_session_id: "source-session",
status: "Review",
})
.await
.expect("failed to fork session snapshot");
let session_rows = database
.sessions()
.load_sessions()
.await
.expect("failed to load sessions");
let source_row = session_rows
.iter()
.find(|session_row| session_row.id == "source-session")
.expect("missing source session row");
let fork_row = session_rows
.iter()
.find(|session_row| session_row.id == "fork-session")
.expect("missing forked session row");
let fork_reset_row = load_fork_reset_row(&pool, "fork-session").await;
let fork_review_request = database
.reviews()
.load_session_review_request("fork-session")
.await
.expect("failed to load fork review request");
assert_source_reset_state(
source_row,
&source_reset_row,
source_review_request.as_ref(),
);
assert_fork_reset_state(fork_row, &fork_reset_row, fork_review_request.as_ref());
}
#[tokio::test]
async fn test_clear_session_draft_flag_marks_draft_session_live() {
let (database, _pool) = AppRepositories::in_memory_with_pool()
.await
.expect("db should open");
let project_id = database
.projects()
.upsert_project("/tmp/project", None)
.await
.expect("failed to upsert project");
database
.sessions()
.insert_draft_session("draft-session", "gpt-5.6-sol", "main", "Draft", project_id)
.await
.expect("failed to insert draft session");
database
.sessions()
.clear_session_draft_flag("draft-session")
.await
.expect("failed to clear session draft flag");
let session_row = database
.sessions()
.load_sessions()
.await
.expect("failed to load sessions")
.into_iter()
.find(|session_row| session_row.id == "draft-session")
.expect("missing draft session row");
assert!(!session_row.is_draft);
}
#[test]
fn test_session_join_row_ignores_partial_review_request_columns() {
let mut session_join_row = SessionJoinRow::fixture_for_test();
session_join_row.review_request_last_refreshed_at = None;
let session_row = session_join_row.into_session_row();
assert_eq!(session_row.id, "session-a");
assert_eq!(session_row.project_id, Some(7));
assert_eq!(
session_row.parent_session_id.as_deref(),
Some("parent-session")
);
assert_eq!(session_row.status, "Review");
assert_eq!(session_row.added_lines, 14);
assert_eq!(session_row.deleted_lines, 6);
assert_eq!(session_row.review_request, None);
}
#[test]
fn test_session_join_row_maps_review_request_columns() {
let session_join_row = SessionJoinRow::fixture_for_test();
let session_row = session_join_row.into_session_row();
assert_eq!(session_row.id, "session-a");
assert_eq!(session_row.added_lines, 14);
assert_eq!(session_row.deleted_lines, 6);
assert_eq!(session_row.project_id, Some(7));
assert_eq!(session_row.personality_id.as_deref(), Some("reviewer"));
assert_eq!(
session_row.parent_session_id.as_deref(),
Some("parent-session")
);
assert_eq!(
session_row.published_upstream_ref.as_deref(),
Some("origin/session-a")
);
assert_eq!(session_row.questions.as_deref(), Some("Question text"));
assert_eq!(session_row.summary.as_deref(), Some("Summary text"));
assert_eq!(session_row.title.as_deref(), Some("Review session"));
assert_eq!(
session_row.review_request,
Some(expected_review_request_row())
);
}
}