Compare commits
3 Commits
dbc676332a
...
4ba2f7af74
| Author | SHA1 | Date | |
|---|---|---|---|
| 4ba2f7af74 | |||
| 08b3cd8fb7 | |||
| 5fd1698d47 |
@@ -12,6 +12,8 @@ harness-mcp = { workspace = true }
|
||||
harness-lsp = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
tokio-util = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
dirs = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
|
||||
@@ -14,8 +14,12 @@ use async_trait::async_trait;
|
||||
use harness_core::agent::{AgentDef, AgentRegistry};
|
||||
use harness_core::config::{self, Config, ConfigError};
|
||||
use harness_core::engine::jobs::{JobBoard, JobState, LaunchSpec};
|
||||
use harness_core::engine::{run_session, RunConfig, StepContext};
|
||||
use harness_core::engine::{run_session, Compactor, RunConfig, StepContext};
|
||||
use harness_core::event::{AppEvent, EventBus, RunOutcome};
|
||||
use harness_core::llm::{
|
||||
Initiator, LlmEvent, LlmRequest, Provider, ProviderError, Role as WireRole, WireContent,
|
||||
WireMessage,
|
||||
};
|
||||
use harness_core::lsp::DiagnosticsSource;
|
||||
use harness_core::permission::{PermissionReply, PermissionService, Rule, Ruleset};
|
||||
use harness_core::store::{Store, StoreError};
|
||||
@@ -164,6 +168,12 @@ struct EngineInner {
|
||||
skills_prompt: Option<String>,
|
||||
/// Slash commands by name, expanded into the user message by the TUI input layer.
|
||||
commands: HashMap<String, config::CommandDef>,
|
||||
/// Auto-compaction backend (small model). `None` when no `small_model` is configured or
|
||||
/// its provider isn't registered — compaction is then disabled.
|
||||
compactor: Option<Arc<dyn Compactor>>,
|
||||
/// Resolved `small_model` (provider + model id) for session-title generation. `None`
|
||||
/// disables titling (the session keeps its empty title / prompt fallback).
|
||||
small_model: Option<(Arc<dyn Provider>, String)>,
|
||||
cwd: PathBuf,
|
||||
data_dir: PathBuf,
|
||||
runs: Mutex<HashMap<SessionId, RunHandle>>,
|
||||
@@ -255,6 +265,22 @@ impl EngineHandle {
|
||||
// `init` warms the cache in the background for the next launch.
|
||||
let catalog = ModelCatalog::load_cached_or_baked(&ModelCatalog::default_cache_path());
|
||||
|
||||
// The configured `small_model` (provider/model), resolved to a provider if registered.
|
||||
// Powers both auto-compaction and session-title generation; `None` disables both.
|
||||
let small_model: Option<(Arc<dyn Provider>, String)> = config
|
||||
.small_model
|
||||
.as_deref()
|
||||
.and_then(|m| m.split_once('/'))
|
||||
.and_then(|(provider_id, model_id)| {
|
||||
providers
|
||||
.get(provider_id)
|
||||
.map(|provider| (provider, model_id.to_string()))
|
||||
});
|
||||
let compactor: Option<Arc<dyn Compactor>> =
|
||||
small_model.clone().map(|(provider, model_id)| {
|
||||
Arc::new(SmallModelCompactor::new(provider, model_id)) as Arc<dyn Compactor>
|
||||
});
|
||||
|
||||
// LSP pool: built-ins present on PATH, plus any config-declared servers.
|
||||
let lsp_servers: Vec<harness_lsp::ServerConfig> = config
|
||||
.lsp
|
||||
@@ -282,6 +308,8 @@ impl EngineHandle {
|
||||
diagnostics,
|
||||
skills_prompt,
|
||||
commands,
|
||||
compactor,
|
||||
small_model,
|
||||
cwd,
|
||||
data_dir,
|
||||
runs: Mutex::new(HashMap::new()),
|
||||
@@ -366,6 +394,12 @@ impl EngineHandle {
|
||||
}
|
||||
}
|
||||
|
||||
// Crash-resume: a previous run may have died mid-turn (process killed) leaving an
|
||||
// assistant message with no finish/error and tool parts stuck Running/Pending. Repair
|
||||
// that before appending the new turn so history is well-formed and the model isn't
|
||||
// shown a dangling tool call.
|
||||
self.inner.repair_unfinished(&session_id).await?;
|
||||
|
||||
let now = now_ms();
|
||||
let user_message = Message::new_user(session_id.clone(), now);
|
||||
self.inner
|
||||
@@ -379,6 +413,9 @@ impl EngineHandle {
|
||||
message: user_message.clone(),
|
||||
});
|
||||
|
||||
// Fire-and-forget session-title generation from the first user turn's text.
|
||||
self.inner.spawn_title_generation(&session_id, &text);
|
||||
|
||||
let user_part = Part {
|
||||
id: PartId::new(),
|
||||
message_id: user_message.id.clone(),
|
||||
@@ -427,6 +464,7 @@ impl EngineHandle {
|
||||
job_board: Some(board),
|
||||
context_reporter: None, // root session has no parent board to report to
|
||||
diagnostics: self.inner.diagnostics.clone(),
|
||||
compactor: self.inner.compactor.clone(),
|
||||
};
|
||||
let run_config = RunConfig {
|
||||
agent_name: agent.name.clone(),
|
||||
@@ -440,6 +478,12 @@ impl EngineHandle {
|
||||
reminder_turn_start: self.inner.reminder("turn_start"),
|
||||
reminder_after_file_tool: self.inner.reminder("after_file_tool"),
|
||||
skills_prompt: self.inner.skills_prompt.clone(),
|
||||
context_limit: self
|
||||
.inner
|
||||
.catalog
|
||||
.get(provider_id, model_id)
|
||||
.map(|i| i.context_limit)
|
||||
.unwrap_or(0),
|
||||
};
|
||||
|
||||
// Reserve the run slot *before* spawning. If we inserted after spawning, a run that
|
||||
@@ -505,6 +549,13 @@ impl EngineHandle {
|
||||
Ok(self.inner.store.session(session_id).await?)
|
||||
}
|
||||
|
||||
/// Repairs a crashed-mid-run session (unfinished assistant → aborted, in-flight tool parts
|
||||
/// → interrupted). Idempotent; safe to call whenever a session is loaded/resumed.
|
||||
pub async fn repair_session(&self, session_id: &SessionId) -> Result<(), AppError> {
|
||||
self.inner.repair_unfinished(session_id).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Background jobs spawned by `session_id` (the parent), for the TUI jobs pane / drill-in.
|
||||
pub async fn jobs(
|
||||
&self,
|
||||
@@ -557,6 +608,116 @@ impl EngineInner {
|
||||
.expect("EngineInner self-reference set in new()")
|
||||
}
|
||||
|
||||
/// Repairs a session that crashed mid-run: any assistant message with neither a finish
|
||||
/// reason nor an error is marked `Aborted`, and its in-flight tool parts (Running/Pending)
|
||||
/// become `Error { "interrupted" }`. Idempotent — a clean session is left untouched.
|
||||
async fn repair_unfinished(&self, session_id: &SessionId) -> Result<(), StoreError> {
|
||||
let messages = self.store.messages(session_id.clone()).await?;
|
||||
for message in messages {
|
||||
let crashed = message.role == harness_core::types::Role::Assistant
|
||||
&& message.finished.is_none()
|
||||
&& message.error.is_none();
|
||||
if !crashed {
|
||||
continue;
|
||||
}
|
||||
|
||||
for mut part in self.store.parts(message.id.clone()).await? {
|
||||
let interrupted = match &part.body {
|
||||
PartBody::Tool { state, .. } => matches!(
|
||||
state,
|
||||
harness_core::types::ToolState::Running { .. }
|
||||
| harness_core::types::ToolState::Pending { .. }
|
||||
),
|
||||
_ => false,
|
||||
};
|
||||
if !interrupted {
|
||||
continue;
|
||||
}
|
||||
if let PartBody::Tool {
|
||||
call_id,
|
||||
name,
|
||||
state,
|
||||
} = part.body
|
||||
{
|
||||
let input = match &state {
|
||||
harness_core::types::ToolState::Running { input, .. } => input.clone(),
|
||||
_ => serde_json::Value::Null,
|
||||
};
|
||||
part.body = PartBody::Tool {
|
||||
call_id,
|
||||
name,
|
||||
state: harness_core::types::ToolState::Error {
|
||||
input,
|
||||
error: "interrupted".to_string(),
|
||||
},
|
||||
};
|
||||
self.store.upsert_part(part.clone()).await?;
|
||||
self.bus.publish(AppEvent::PartUpdated { part });
|
||||
}
|
||||
}
|
||||
|
||||
let mut repaired = message;
|
||||
repaired.error = Some(harness_core::types::MessageError::Aborted);
|
||||
self.store.upsert_message(repaired.clone()).await?;
|
||||
self.bus
|
||||
.publish(AppEvent::MessageUpdated { message: repaired });
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Fires a background task that generates a short session title from the first user turn.
|
||||
/// No-op when no `small_model` is available; the title is only generated once (the task
|
||||
/// re-checks that the session is still untitled before writing).
|
||||
fn spawn_title_generation(&self, session_id: &SessionId, user_text: &str) {
|
||||
let Some((provider, model_id)) = self.small_model.clone() else {
|
||||
return;
|
||||
};
|
||||
let inner = self.arc();
|
||||
let session_id = session_id.clone();
|
||||
let user_text = user_text.to_string();
|
||||
tokio::spawn(async move {
|
||||
inner
|
||||
.generate_title(provider, model_id, session_id, user_text)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
|
||||
/// Generates and stores a session title for a still-untitled root session.
|
||||
async fn generate_title(
|
||||
&self,
|
||||
provider: Arc<dyn Provider>,
|
||||
model_id: String,
|
||||
session_id: SessionId,
|
||||
user_text: String,
|
||||
) {
|
||||
// Only title a root session that hasn't been titled yet.
|
||||
match self.store.session(session_id.clone()).await {
|
||||
Ok(Some(s)) if s.depth == 0 && s.title.is_empty() => {}
|
||||
_ => return,
|
||||
}
|
||||
let raw = match oneshot_text(&provider, &model_id, TITLE_SYSTEM_PROMPT, &user_text).await {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
tracing::debug!(error = %e, "session-title generation failed");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let title = sanitize_title(&raw);
|
||||
if title.is_empty() {
|
||||
return;
|
||||
}
|
||||
if let Ok(Some(mut session)) = self.store.session(session_id.clone()).await {
|
||||
if !session.title.is_empty() {
|
||||
return; // titled concurrently; don't clobber
|
||||
}
|
||||
session.title = title;
|
||||
session.updated_at = now_ms();
|
||||
if self.store.upsert_session(session.clone()).await.is_ok() {
|
||||
self.bus.publish(AppEvent::SessionUpdated { session });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolves an agent definition by name, falling back to a generic assistant so an
|
||||
/// unknown/misconfigured agent still runs rather than failing the session.
|
||||
fn agent_def(&self, name: &str) -> AgentDef {
|
||||
@@ -734,6 +895,7 @@ impl EngineInner {
|
||||
job_board: Some(board),
|
||||
context_reporter: reporter,
|
||||
diagnostics: self.diagnostics.clone(),
|
||||
compactor: self.compactor.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -754,6 +916,101 @@ impl ContextReporter for BoardReporter {
|
||||
}
|
||||
}
|
||||
|
||||
/// System prompt for session-title generation.
|
||||
const TITLE_SYSTEM_PROMPT: &str = "Generate a concise title (3 to 6 words) for a coding \
|
||||
session, summarizing what the user is asking for. No quotes, no trailing punctuation, no \
|
||||
preamble. Output only the title.";
|
||||
|
||||
/// Cap on a generated title's length, in characters.
|
||||
const MAX_TITLE_LEN: usize = 60;
|
||||
|
||||
/// Cleans a model-generated title: first non-empty line, stripped of surrounding quotes, and
|
||||
/// truncated to [`MAX_TITLE_LEN`] characters (on a char boundary).
|
||||
fn sanitize_title(raw: &str) -> String {
|
||||
let line = raw
|
||||
.lines()
|
||||
.map(str::trim)
|
||||
.find(|l| !l.is_empty())
|
||||
.unwrap_or("");
|
||||
let line = line
|
||||
.trim_matches(|c| c == '"' || c == '\'' || c == '`')
|
||||
.trim();
|
||||
line.chars().take(MAX_TITLE_LEN).collect()
|
||||
}
|
||||
|
||||
/// One-shot text completion: sends `user_text` under `system` to the model and concatenates the
|
||||
/// streamed text. Used for session titles (compaction uses the message-array form directly).
|
||||
async fn oneshot_text(
|
||||
provider: &Arc<dyn Provider>,
|
||||
model_id: &str,
|
||||
system: &str,
|
||||
user_text: &str,
|
||||
) -> Result<String, ProviderError> {
|
||||
use futures::StreamExt;
|
||||
|
||||
let req = LlmRequest {
|
||||
model: model_id.to_string(),
|
||||
system: vec![system.to_string()],
|
||||
messages: vec![WireMessage {
|
||||
role: WireRole::User,
|
||||
content: vec![WireContent::Text {
|
||||
text: user_text.to_string(),
|
||||
}],
|
||||
}],
|
||||
tools: Vec::new(),
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
reasoning: None,
|
||||
initiator: Initiator::Agent,
|
||||
};
|
||||
let mut stream = provider.stream(req, CancellationToken::new()).await?;
|
||||
let mut out = String::new();
|
||||
while let Some(event) = stream.next().await {
|
||||
if let LlmEvent::TextDelta { text, .. } = event? {
|
||||
out.push_str(&text);
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Auto-compaction backend: summarizes a conversation with the configured `small_model`.
|
||||
struct SmallModelCompactor {
|
||||
provider: Arc<dyn Provider>,
|
||||
model_id: String,
|
||||
}
|
||||
|
||||
impl SmallModelCompactor {
|
||||
fn new(provider: Arc<dyn Provider>, model_id: String) -> Self {
|
||||
Self { provider, model_id }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Compactor for SmallModelCompactor {
|
||||
async fn summarize(&self, messages: &[WireMessage]) -> Result<String, ProviderError> {
|
||||
use futures::StreamExt;
|
||||
|
||||
let req = LlmRequest {
|
||||
model: self.model_id.clone(),
|
||||
system: vec![harness_core::engine::compact::SUMMARY_SYSTEM_PROMPT.to_string()],
|
||||
messages: messages.to_vec(),
|
||||
tools: Vec::new(),
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
reasoning: None,
|
||||
initiator: Initiator::Agent,
|
||||
};
|
||||
let mut stream = self.provider.stream(req, CancellationToken::new()).await?;
|
||||
let mut summary = String::new();
|
||||
while let Some(event) = stream.next().await {
|
||||
if let LlmEvent::TextDelta { text, .. } = event? {
|
||||
summary.push_str(&text);
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SubagentSpawner for EngineInner {
|
||||
async fn spawn(&self, req: SpawnRequest) -> Result<SpawnOutcome, SpawnError> {
|
||||
@@ -853,6 +1110,14 @@ impl SubagentSpawner for EngineInner {
|
||||
reminder_turn_start: self.reminder("turn_start"),
|
||||
reminder_after_file_tool: self.reminder("after_file_tool"),
|
||||
skills_prompt: self.skills_prompt.clone(),
|
||||
context_limit: self
|
||||
.catalog
|
||||
.get(
|
||||
&child_session.model.provider_id,
|
||||
&child_session.model.model_id,
|
||||
)
|
||||
.map(|i| i.context_limit)
|
||||
.unwrap_or(0),
|
||||
};
|
||||
|
||||
// Register the launch on the parent board (also makes foreground results reusable).
|
||||
@@ -1443,6 +1708,108 @@ mod tests {
|
||||
assert_eq!(stored.agent, "orchestrator");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn repair_marks_crashed_assistant_aborted_and_interrupts_tool_parts() {
|
||||
use harness_core::types::{MessageError, PartId, Role, ToolState};
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let engine = mock_engine(dir.path().to_path_buf());
|
||||
let session = engine
|
||||
.new_session("orchestrator", "mock/mock-model")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// A crashed assistant turn: no finish, no error, with a tool call left Running.
|
||||
let model = ModelRef::new("mock", "mock-model");
|
||||
let crashed = Message::new_assistant(session.clone(), model, "orchestrator", now_ms());
|
||||
engine
|
||||
.inner
|
||||
.store
|
||||
.upsert_message(crashed.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
let part = Part {
|
||||
id: PartId::new(),
|
||||
message_id: crashed.id.clone(),
|
||||
session_id: session.clone(),
|
||||
idx: 0,
|
||||
body: PartBody::Tool {
|
||||
call_id: "c1".into(),
|
||||
name: "bash".into(),
|
||||
state: ToolState::Running {
|
||||
input: serde_json::json!({"command": "sleep 100"}),
|
||||
title: None,
|
||||
metadata: serde_json::Value::Null,
|
||||
},
|
||||
},
|
||||
};
|
||||
engine.inner.store.upsert_part(part).await.unwrap();
|
||||
|
||||
engine.inner.repair_unfinished(&session).await.unwrap();
|
||||
|
||||
let messages = engine.inner.store.messages(session.clone()).await.unwrap();
|
||||
let repaired = messages.iter().find(|m| m.id == crashed.id).unwrap();
|
||||
assert_eq!(repaired.role, Role::Assistant);
|
||||
assert!(matches!(repaired.error, Some(MessageError::Aborted)));
|
||||
|
||||
let parts = engine.inner.store.parts(crashed.id.clone()).await.unwrap();
|
||||
match &parts[0].body {
|
||||
PartBody::Tool { state, .. } => match state {
|
||||
ToolState::Error { error, .. } => assert_eq!(error, "interrupted"),
|
||||
other => panic!("expected Error tool state, got {other:?}"),
|
||||
},
|
||||
other => panic!("expected Tool part, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn first_prompt_generates_a_session_title() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut providers = ProviderRegistry::new();
|
||||
providers.register(Arc::new(MockProvider));
|
||||
let config = Config {
|
||||
small_model: Some("mock/mock-model".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let engine = EngineHandle::build(
|
||||
dir.path().to_path_buf(),
|
||||
Store::open_in_memory().unwrap(),
|
||||
config,
|
||||
providers,
|
||||
Vec::new(),
|
||||
)
|
||||
.unwrap();
|
||||
spawn_auto_approve_task(&engine);
|
||||
let session = engine
|
||||
.new_session("orchestrator", "mock/mock-model")
|
||||
.await
|
||||
.unwrap();
|
||||
let mut rx = engine.bus().subscribe();
|
||||
engine
|
||||
.prompt(
|
||||
session.clone(),
|
||||
"help me refactor the auth module".into(),
|
||||
"mock/mock-model",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// The title-generation task (small model = the mock) publishes a SessionUpdated whose
|
||||
// title is the mock's canned text.
|
||||
let title = tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||||
loop {
|
||||
if let Ok(AppEvent::SessionUpdated { session: s }) = rx.recv().await {
|
||||
if s.id == session && !s.title.is_empty() {
|
||||
return s.title;
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("a titled SessionUpdated within the timeout");
|
||||
assert_eq!(title, "explored the code");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_prompt_errors_on_unregistered_provider() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
//! Auto-compaction (M6). When a session's context approaches the model's window the loop
|
||||
//! summarizes the conversation so far via a small model, writes a `Compaction` marker, and
|
||||
//! continues — subsequent requests replace the summarized history with the summary. See
|
||||
//! `docs/02-engine.md`.
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::event::AppEvent;
|
||||
use crate::llm::{ProviderError, WireMessage};
|
||||
use crate::types::{Message, MessageId, Part, PartBody, PartId, SessionId};
|
||||
|
||||
use super::processor::StepContext;
|
||||
|
||||
/// Produces a compact recap of a conversation. Implemented by the composition root over the
|
||||
/// configured `small_model`; absent in headless/test contexts (compaction then disabled).
|
||||
#[async_trait]
|
||||
pub trait Compactor: Send + Sync {
|
||||
async fn summarize(&self, messages: &[WireMessage]) -> Result<String, ProviderError>;
|
||||
}
|
||||
|
||||
/// System prompt handed to the small model to summarize the conversation for continuation.
|
||||
pub const SUMMARY_SYSTEM_PROMPT: &str = "You are compacting a long coding-assistant \
|
||||
conversation so it can continue within a smaller context window. Write a dense, factual \
|
||||
summary that preserves: the user's goal and constraints, decisions made and why, files \
|
||||
and symbols touched, commands run and their results, and the exact next step in progress. \
|
||||
Omit pleasantries. Output only the summary.";
|
||||
|
||||
/// Frames a raw summary as the synthetic user turn the model sees after compaction.
|
||||
pub fn frame_summary(summary: &str) -> String {
|
||||
format!(
|
||||
"The earlier part of this conversation was summarized to save context. \
|
||||
Continue from this summary:\n\n{summary}"
|
||||
)
|
||||
}
|
||||
|
||||
/// Persists a compaction: a new user message carrying the `Compaction` marker (used to cut
|
||||
/// history on the next load) plus a synthetic text part holding the framed summary (what the
|
||||
/// model actually reads). Returns the new message id.
|
||||
pub async fn write_compaction(
|
||||
ctx: &StepContext,
|
||||
session_id: &SessionId,
|
||||
replaces_up_to: MessageId,
|
||||
summary: String,
|
||||
now: i64,
|
||||
) -> Result<MessageId, ProviderError> {
|
||||
let message = Message::new_user(session_id.clone(), now);
|
||||
let message_id = message.id.clone();
|
||||
let store_err = |e: crate::store::StoreError| ProviderError::Decode(e.to_string());
|
||||
|
||||
ctx.store
|
||||
.upsert_message(message.clone())
|
||||
.await
|
||||
.map_err(store_err)?;
|
||||
ctx.bus.publish(AppEvent::MessageCreated {
|
||||
message: message.clone(),
|
||||
});
|
||||
|
||||
let marker = Part {
|
||||
id: PartId::new(),
|
||||
message_id: message_id.clone(),
|
||||
session_id: session_id.clone(),
|
||||
idx: 0,
|
||||
body: PartBody::Compaction {
|
||||
replaces_up_to,
|
||||
summary: summary.clone(),
|
||||
},
|
||||
};
|
||||
ctx.store.upsert_part(marker).await.map_err(store_err)?;
|
||||
|
||||
let text = Part {
|
||||
id: PartId::new(),
|
||||
message_id: message_id.clone(),
|
||||
session_id: session_id.clone(),
|
||||
idx: 1,
|
||||
body: PartBody::Text {
|
||||
text: frame_summary(&summary),
|
||||
synthetic: true,
|
||||
},
|
||||
};
|
||||
ctx.store
|
||||
.upsert_part(text.clone())
|
||||
.await
|
||||
.map_err(store_err)?;
|
||||
ctx.bus.publish(AppEvent::PartUpdated { part: text });
|
||||
|
||||
Ok(message_id)
|
||||
}
|
||||
|
||||
/// Index of the last message carrying a `Compaction` part, given each message's parts. History
|
||||
/// before it is dropped when building the request; `None` means no compaction yet.
|
||||
pub fn last_compaction_index(parts_per_message: &[Vec<Part>]) -> Option<usize> {
|
||||
parts_per_message.iter().rposition(|parts| {
|
||||
parts
|
||||
.iter()
|
||||
.any(|p| matches!(p.body, PartBody::Compaction { .. }))
|
||||
})
|
||||
}
|
||||
|
||||
/// Whether accumulated `used` tokens have crossed the compaction threshold (90% of the
|
||||
/// model's context window). `context_limit == 0` (unknown) disables the trigger.
|
||||
pub fn over_threshold(used: u64, context_limit: u64) -> bool {
|
||||
context_limit > 0 && used.saturating_mul(10) > context_limit.saturating_mul(9)
|
||||
}
|
||||
|
||||
/// Effective context occupancy for a step's usage: prompt (incl. cache) plus generated output.
|
||||
pub fn tokens_used(usage: &crate::types::TokenUsage) -> u64 {
|
||||
usage.input + usage.cache_read + usage.cache_write + usage.output
|
||||
}
|
||||
|
||||
/// True if a message list has real history to compact before `cut_start` (avoids a useless
|
||||
/// compaction that would summarize nothing and could loop).
|
||||
pub fn has_history_to_compact(messages: &[Message], cut_start: usize) -> bool {
|
||||
messages.len().saturating_sub(cut_start) > 1
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn over_threshold_respects_ninety_percent_and_unknown_limit() {
|
||||
assert!(over_threshold(91, 100));
|
||||
assert!(!over_threshold(90, 100)); // exactly 90% is not yet over
|
||||
assert!(!over_threshold(50, 100));
|
||||
assert!(!over_threshold(1_000_000, 0)); // unknown limit disables the trigger
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tokens_used_sums_prompt_cache_and_output() {
|
||||
let usage = crate::types::TokenUsage {
|
||||
input: 10,
|
||||
output: 5,
|
||||
reasoning: 0,
|
||||
cache_read: 3,
|
||||
cache_write: 2,
|
||||
};
|
||||
assert_eq!(tokens_used(&usage), 20);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn frame_summary_wraps_text() {
|
||||
let framed = frame_summary("did X");
|
||||
assert!(framed.contains("did X"));
|
||||
assert!(framed.contains("summarized"));
|
||||
}
|
||||
}
|
||||
@@ -30,6 +30,9 @@ pub struct RunConfig {
|
||||
pub reminder_after_file_tool: Option<String>,
|
||||
/// Pre-rendered "## Skills" system block advertising loadable skills. `None` = no skills.
|
||||
pub skills_prompt: Option<String>,
|
||||
/// The model's context window in tokens, for the auto-compaction trigger. `0` = unknown
|
||||
/// (disables the threshold trigger; a hard `ContextOverflow` still compacts).
|
||||
pub context_limit: u64,
|
||||
}
|
||||
|
||||
/// Adds a step's usage/cost onto the persisted session and republishes it. Cost accounting is
|
||||
@@ -157,6 +160,54 @@ fn convert_message(message: &Message, parts: &[Part]) -> Vec<WireMessage> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads the session's messages, applies the compaction cut (drop everything before the last
|
||||
/// `Compaction` marker), and converts the surviving window to wire messages. Returns the full
|
||||
/// message list (for the "is there history to compact" check) alongside the wire window and
|
||||
/// the cut start index.
|
||||
async fn load_wire(
|
||||
ctx: &StepContext,
|
||||
) -> Result<(Vec<Message>, Vec<WireMessage>, usize), ProviderError> {
|
||||
let messages = ctx
|
||||
.store
|
||||
.messages(ctx.session_id.clone())
|
||||
.await
|
||||
.map_err(|e| ProviderError::Decode(e.to_string()))?;
|
||||
let mut parts_per_message = Vec::with_capacity(messages.len());
|
||||
for message in &messages {
|
||||
let parts = ctx
|
||||
.store
|
||||
.parts(message.id.clone())
|
||||
.await
|
||||
.map_err(|e| ProviderError::Decode(e.to_string()))?;
|
||||
parts_per_message.push(parts);
|
||||
}
|
||||
let cut_start = super::compact::last_compaction_index(&parts_per_message).unwrap_or(0);
|
||||
let mut wire = Vec::new();
|
||||
for (message, parts) in messages.iter().zip(&parts_per_message).skip(cut_start) {
|
||||
wire.extend(convert_message(message, parts));
|
||||
}
|
||||
Ok((messages, wire, cut_start))
|
||||
}
|
||||
|
||||
/// Summarizes the current context window and persists a compaction marker. Returns `true` if a
|
||||
/// compaction was actually written (there was history to summarize), `false` if there was
|
||||
/// nothing to compact. Requires `ctx.compactor` to be set.
|
||||
async fn do_compaction(ctx: &StepContext, now: i64) -> Result<bool, ProviderError> {
|
||||
let Some(compactor) = ctx.compactor.clone() else {
|
||||
return Ok(false);
|
||||
};
|
||||
let (messages, wire, cut_start) = load_wire(ctx).await?;
|
||||
if !super::compact::has_history_to_compact(&messages, cut_start) {
|
||||
return Ok(false);
|
||||
}
|
||||
let Some(last) = messages.last() else {
|
||||
return Ok(false);
|
||||
};
|
||||
let summary = compactor.summarize(&wire).await?;
|
||||
super::compact::write_compaction(ctx, &ctx.session_id, last.id.clone(), summary, now).await?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
/// The outer loop: one call per session turn until the model stops asking for tool calls.
|
||||
/// `now_fn` supplies `created_at`/`StepContext::now` stamps (kept out of the loop body so
|
||||
/// tests can drive deterministic timestamps).
|
||||
@@ -187,26 +238,14 @@ pub async fn run_session(
|
||||
steps += 1;
|
||||
ctx.now = now_fn();
|
||||
|
||||
let messages = match ctx.store.messages(ctx.session_id.clone()).await {
|
||||
Ok(m) => m,
|
||||
let mut wire_messages = match load_wire(&ctx).await {
|
||||
Ok((_, wire, _)) => wire,
|
||||
Err(e) => {
|
||||
return RunOutcome::Errored {
|
||||
message: e.to_string(),
|
||||
}
|
||||
}
|
||||
};
|
||||
let mut wire_messages = Vec::new();
|
||||
for message in &messages {
|
||||
let parts = match ctx.store.parts(message.id.clone()).await {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
return RunOutcome::Errored {
|
||||
message: e.to_string(),
|
||||
}
|
||||
}
|
||||
};
|
||||
wire_messages.extend(convert_message(message, &parts));
|
||||
}
|
||||
|
||||
// Collect synthetic (non-persisted) blocks to append to the last user message this
|
||||
// turn: the optional turn-start reminder, the job board, and — if the previous step
|
||||
@@ -312,15 +351,49 @@ pub async fn run_session(
|
||||
}
|
||||
}
|
||||
}
|
||||
// Proactive compaction: if the context has crossed ~90% of the window and a
|
||||
// compactor is wired in, summarize before the next request. Best-effort — a
|
||||
// compaction failure just means we continue with the full history.
|
||||
let near_limit = ctx.compactor.is_some()
|
||||
&& super::compact::over_threshold(
|
||||
super::compact::tokens_used(&outcome.usage),
|
||||
run_config.context_limit,
|
||||
);
|
||||
match outcome.result {
|
||||
StepResult::Continue => continue,
|
||||
StepResult::Continue | StepResult::Compact => {
|
||||
if near_limit || matches!(outcome.result, StepResult::Compact) {
|
||||
if let Err(e) = do_compaction(&ctx, now_fn()).await {
|
||||
tracing::warn!(error = %e, "compaction failed; continuing with full history");
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
StepResult::Stop => return RunOutcome::Stopped,
|
||||
StepResult::Compact => return RunOutcome::Stopped, // stub until M6
|
||||
}
|
||||
}
|
||||
Err(step_err) if matches!(step_err.source, ProviderError::Cancelled) => {
|
||||
return RunOutcome::Aborted;
|
||||
}
|
||||
// A hard context-window overflow is recoverable when a compactor is available:
|
||||
// summarize and retry. If there was nothing left to compact, surface the error.
|
||||
Err(step_err)
|
||||
if matches!(step_err.source, ProviderError::ContextOverflow)
|
||||
&& ctx.compactor.is_some() =>
|
||||
{
|
||||
match do_compaction(&ctx, now_fn()).await {
|
||||
Ok(true) => continue,
|
||||
Ok(false) => {
|
||||
return RunOutcome::Errored {
|
||||
message: step_err.source.to_string(),
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
return RunOutcome::Errored {
|
||||
message: e.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(step_err) => {
|
||||
return RunOutcome::Errored {
|
||||
message: step_err.source.to_string(),
|
||||
@@ -438,6 +511,7 @@ mod tests {
|
||||
job_board: None,
|
||||
context_reporter: None,
|
||||
diagnostics: None,
|
||||
compactor: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -523,6 +597,7 @@ mod tests {
|
||||
reminder_turn_start: None,
|
||||
reminder_after_file_tool: None,
|
||||
skills_prompt: None,
|
||||
context_limit: 0,
|
||||
};
|
||||
|
||||
let outcome = run_session(Arc::new(provider), ctx, &run_config, || 2).await;
|
||||
@@ -586,6 +661,207 @@ mod tests {
|
||||
assert_eq!(final_text, "The file contains: mock file content");
|
||||
}
|
||||
|
||||
/// Records its calls and returns a fixed summary, so a test can assert compaction ran.
|
||||
struct MockCompactor {
|
||||
calls: Arc<std::sync::atomic::AtomicUsize>,
|
||||
summary: String,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl crate::engine::compact::Compactor for MockCompactor {
|
||||
async fn summarize(&self, _messages: &[WireMessage]) -> Result<String, ProviderError> {
|
||||
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
Ok(self.summary.clone())
|
||||
}
|
||||
}
|
||||
|
||||
async fn seed_session_with_user_text(store: &Store, model: &ModelRef, text: &str) -> SessionId {
|
||||
let session = Session::new_root("orchestrator", model.clone(), 1);
|
||||
let session_id = session.id.clone();
|
||||
store.upsert_session(session).await.unwrap();
|
||||
let user_message = Message::new_user(session_id.clone(), 1);
|
||||
store.upsert_message(user_message.clone()).await.unwrap();
|
||||
store
|
||||
.upsert_part(Part {
|
||||
id: crate::types::PartId::new(),
|
||||
message_id: user_message.id.clone(),
|
||||
session_id: session_id.clone(),
|
||||
idx: 0,
|
||||
body: PartBody::Text {
|
||||
text: text.into(),
|
||||
synthetic: false,
|
||||
},
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
session_id
|
||||
}
|
||||
|
||||
fn run_config_with_limit(model: ModelRef, context_limit: u64) -> RunConfig {
|
||||
RunConfig {
|
||||
agent_name: "orchestrator".into(),
|
||||
agent_prompt: "You are a helpful assistant.".into(),
|
||||
model,
|
||||
temperature: None,
|
||||
max_steps: 10,
|
||||
instructions: Vec::new(),
|
||||
cost: None,
|
||||
inject_job_board: false,
|
||||
reminder_turn_start: None,
|
||||
reminder_after_file_tool: None,
|
||||
skills_prompt: None,
|
||||
context_limit,
|
||||
}
|
||||
}
|
||||
|
||||
async fn compaction_parts(
|
||||
store: &Store,
|
||||
session_id: &SessionId,
|
||||
) -> Vec<(crate::types::MessageId, String)> {
|
||||
let messages = store.messages(session_id.clone()).await.unwrap();
|
||||
let mut found = Vec::new();
|
||||
for message in messages {
|
||||
for part in store.parts(message.id.clone()).await.unwrap() {
|
||||
if let PartBody::Compaction {
|
||||
replaces_up_to,
|
||||
summary,
|
||||
} = part.body
|
||||
{
|
||||
found.push((replaces_up_to, summary));
|
||||
}
|
||||
}
|
||||
}
|
||||
found
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn crossing_the_threshold_compacts_then_continues() {
|
||||
let store = Store::open_in_memory().unwrap();
|
||||
let bus = EventBus::new();
|
||||
let model = ModelRef::new("mock", "mock-model");
|
||||
let session_id = seed_session_with_user_text(&store, &model, "do a long task").await;
|
||||
|
||||
// Step 1: a tool-less "continue" whose usage (200) blows past 90% of a 100-token window.
|
||||
// Step 2: a final answer, run against the compacted history.
|
||||
let provider = MockProvider::scripted(vec![
|
||||
vec![
|
||||
Ok(LlmEvent::TextStart { id: "t1".into() }),
|
||||
Ok(LlmEvent::TextDelta {
|
||||
id: "t1".into(),
|
||||
text: "working".into(),
|
||||
}),
|
||||
Ok(LlmEvent::TextEnd { id: "t1".into() }),
|
||||
Ok(LlmEvent::Finish {
|
||||
reason: FinishReason::ToolCalls,
|
||||
usage: usage(200, 0),
|
||||
}),
|
||||
],
|
||||
vec![
|
||||
Ok(LlmEvent::TextStart { id: "t2".into() }),
|
||||
Ok(LlmEvent::TextDelta {
|
||||
id: "t2".into(),
|
||||
text: "done".into(),
|
||||
}),
|
||||
Ok(LlmEvent::TextEnd { id: "t2".into() }),
|
||||
Ok(LlmEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
usage: usage(5, 1),
|
||||
}),
|
||||
],
|
||||
]);
|
||||
|
||||
let cwd = tempfile::tempdir().unwrap();
|
||||
let mut ctx = make_ctx(
|
||||
store.clone(),
|
||||
bus,
|
||||
session_id.clone(),
|
||||
cwd.path().to_path_buf(),
|
||||
)
|
||||
.await;
|
||||
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
ctx.compactor = Some(Arc::new(MockCompactor {
|
||||
calls: calls.clone(),
|
||||
summary: "COMPACTED SUMMARY".into(),
|
||||
}));
|
||||
|
||||
let run_config = run_config_with_limit(model, 100);
|
||||
let outcome = run_session(Arc::new(provider), ctx, &run_config, || 2).await;
|
||||
assert!(matches!(outcome, RunOutcome::Stopped));
|
||||
assert_eq!(
|
||||
calls.load(std::sync::atomic::Ordering::SeqCst),
|
||||
1,
|
||||
"compacted once"
|
||||
);
|
||||
|
||||
let parts = compaction_parts(&store, &session_id).await;
|
||||
assert_eq!(parts.len(), 1, "one compaction marker written");
|
||||
assert_eq!(parts[0].1, "COMPACTED SUMMARY");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn context_overflow_recovers_by_compacting() {
|
||||
let store = Store::open_in_memory().unwrap();
|
||||
let bus = EventBus::new();
|
||||
let model = ModelRef::new("mock", "mock-model");
|
||||
let session_id = seed_session_with_user_text(&store, &model, "hello").await;
|
||||
// A prior assistant turn (still "in tool calls") so the loop keeps going into step 1 and
|
||||
// there is real history to compact when the overflow hits.
|
||||
let mut prior =
|
||||
Message::new_assistant(session_id.clone(), model.clone(), "orchestrator", 1);
|
||||
prior.finished = Some(FinishReason::ToolCalls);
|
||||
store.upsert_message(prior.clone()).await.unwrap();
|
||||
store
|
||||
.upsert_part(Part {
|
||||
id: crate::types::PartId::new(),
|
||||
message_id: prior.id.clone(),
|
||||
session_id: session_id.clone(),
|
||||
idx: 0,
|
||||
body: PartBody::Text {
|
||||
text: "earlier work".into(),
|
||||
synthetic: false,
|
||||
},
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Step 1: hard context overflow (no output). Step 2: succeeds post-compaction.
|
||||
let provider = MockProvider::scripted(vec![
|
||||
vec![Err(ProviderError::ContextOverflow)],
|
||||
vec![
|
||||
Ok(LlmEvent::TextStart { id: "t".into() }),
|
||||
Ok(LlmEvent::TextDelta {
|
||||
id: "t".into(),
|
||||
text: "recovered".into(),
|
||||
}),
|
||||
Ok(LlmEvent::TextEnd { id: "t".into() }),
|
||||
Ok(LlmEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
usage: usage(1, 1),
|
||||
}),
|
||||
],
|
||||
]);
|
||||
|
||||
let cwd = tempfile::tempdir().unwrap();
|
||||
let mut ctx = make_ctx(
|
||||
store.clone(),
|
||||
bus,
|
||||
session_id.clone(),
|
||||
cwd.path().to_path_buf(),
|
||||
)
|
||||
.await;
|
||||
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
ctx.compactor = Some(Arc::new(MockCompactor {
|
||||
calls: calls.clone(),
|
||||
summary: "RECAP".into(),
|
||||
}));
|
||||
|
||||
let run_config = run_config_with_limit(model, 0); // threshold off; only overflow triggers
|
||||
let outcome = run_session(Arc::new(provider), ctx, &run_config, || 2).await;
|
||||
assert!(matches!(outcome, RunOutcome::Stopped));
|
||||
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
|
||||
assert_eq!(compaction_parts(&store, &session_id).await.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_error_on_first_event_is_reported_and_run_errors() {
|
||||
let store = Store::open_in_memory().unwrap();
|
||||
@@ -634,6 +910,7 @@ mod tests {
|
||||
reminder_turn_start: None,
|
||||
reminder_after_file_tool: None,
|
||||
skills_prompt: None,
|
||||
context_limit: 0,
|
||||
};
|
||||
|
||||
let outcome = run_session(Arc::new(provider), ctx, &run_config, || 2).await;
|
||||
@@ -742,6 +1019,7 @@ mod tests {
|
||||
reminder_turn_start: None,
|
||||
reminder_after_file_tool: None,
|
||||
skills_prompt: None,
|
||||
context_limit: 0,
|
||||
};
|
||||
|
||||
let provider = std::sync::Arc::new(CapturingProvider {
|
||||
@@ -810,6 +1088,7 @@ mod tests {
|
||||
reminder_turn_start: Some("REMEMBER: stay on task.".into()),
|
||||
reminder_after_file_tool: None,
|
||||
skills_prompt: None,
|
||||
context_limit: 0,
|
||||
};
|
||||
|
||||
let provider = std::sync::Arc::new(CapturingProvider {
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod compact;
|
||||
pub mod doomloop;
|
||||
pub mod jobs;
|
||||
pub mod processor;
|
||||
@@ -6,6 +7,7 @@ pub mod retry;
|
||||
pub mod session_loop;
|
||||
pub mod system;
|
||||
|
||||
pub use compact::Compactor;
|
||||
pub use doomloop::DoomLoopGuard;
|
||||
pub use jobs::{ContextFile, JobBoard, JobRecord, JobState};
|
||||
pub use processor::{process_step, StepContext, StepError, StepOutcome, StepResult};
|
||||
|
||||
@@ -74,6 +74,9 @@ pub struct StepContext {
|
||||
pub context_reporter: Option<Arc<dyn ContextReporter>>,
|
||||
/// Language-server diagnostics source shared by the session's edit/write tool calls.
|
||||
pub diagnostics: Option<Arc<dyn crate::lsp::DiagnosticsSource>>,
|
||||
/// Summarizes history when the context nears the model's window. `None` disables
|
||||
/// auto-compaction (headless/tests, or when no `small_model` is configured).
|
||||
pub compactor: Option<Arc<dyn super::compact::Compactor>>,
|
||||
}
|
||||
|
||||
struct FlushTracker {
|
||||
|
||||
@@ -26,6 +26,9 @@ pub async fn select_session(
|
||||
session: Session,
|
||||
) -> Result<(), harness_app::AppError> {
|
||||
let session_id = session.id.clone();
|
||||
// Repair a session that crashed mid-run before loading it, so the transcript shows the
|
||||
// aborted turn rather than a tool call frozen "running".
|
||||
engine.repair_session(&session_id).await?;
|
||||
let messages = engine.session_messages(session_id.clone()).await?;
|
||||
let mut all_parts = Vec::new();
|
||||
for message in &messages {
|
||||
|
||||
Reference in New Issue
Block a user