Skip to content

Commit 53f570b

Browse files
author
user
committed
feat(session): add compact action to session control tool
Add the compact action to the session control tool family so a caller can compact a session context to reduce its token footprint and surface the applied compaction statistics (tokens/ratio/summary). - Thread ContextCompactionOutcome through the manual compaction task so the finished outcome is returned instead of being discarded (unit). - Add compact_session_with_outcome returning BitFunResult<ContextCompactionOutcome>; compact_session_manually delegates to it for the Desktop compatibility API. - Extend SessionControlAction with a Compact variant and its as_str mapping. - Allow compacting the caller own session (the sole self-mutation exception) while keeping owner/creator/ancestor authorization inline. - Wire the Compact dispatch to the compact engine and return the applied stats. - Extend the tool input schema/description and add compact validation tests. Test: cargo check -p bitfun-core --jobs 4 EXIT 0; cargo check -p bitfun-agent-runtime --jobs 4 EXIT 0; cargo test -p bitfun-agent-runtime --features agent-runtime compact --jobs 4 (5 tests pass); cargo test -p bitfun-core --features agent-runtime session_control_tool --jobs 4 (14 tests pass incl validate_compact_*); cargo test -p bitfun-core --features agent-runtime manual_compaction --jobs 4 (6 tests pass). AI: lightly tested (targeted cargo check + compact/session_control_tool/manual_compaction unit tests).
1 parent bc73ae3 commit 53f570b

3 files changed

Lines changed: 272 additions & 8 deletions

File tree

src/crates/assembly/core/src/agentic/coordination/coordinator.rs

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -789,7 +789,7 @@ struct SessionExecutionLease {
789789

790790
struct ManualCompactionTask {
791791
turn_id: String,
792-
completion: oneshot::Receiver<BitFunResult<()>>,
792+
completion: oneshot::Receiver<BitFunResult<ContextCompactionOutcome>>,
793793
}
794794

795795
struct ManualCompactionControlGuard {
@@ -5451,7 +5451,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet
54515451
remote_exec_port: Option<Arc<dyn RemoteExecPort>>,
54525452
cancellation_token: CancellationToken,
54535453
commit_gate: Arc<ManualCompactionCommitGate>,
5454-
) -> BitFunResult<()> {
5454+
) -> BitFunResult<ContextCompactionOutcome> {
54555455
let manual_workspace_services = Self::build_workspace_services(&manual_workspace).await;
54565456
let manual_execution_context = ExecutionContext {
54575457
session_id: session_id.clone(),
@@ -5516,7 +5516,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet
55165516
&outcome,
55175517
context_window,
55185518
)
5519-
.await
5519+
.await?;
5520+
Ok(outcome)
55205521
}
55215522
Err(err @ BitFunError::Cancelled(_)) => {
55225523
let error_text = err.to_string();
@@ -5584,6 +5585,17 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet
55845585
/// task used by Agent Runtime callers, then await its terminal result for
55855586
/// the existing Desktop compatibility API.
55865587
pub async fn compact_session_manually(&self, session_id: String) -> BitFunResult<()> {
5588+
self.compact_session_with_outcome(session_id)
5589+
.await
5590+
.map(|_| ())
5591+
}
5592+
5593+
/// Compact the active session context and return the compaction outcome
5594+
/// (tokens/ratio/summary) so tool callers can surface the applied result.
5595+
pub async fn compact_session_with_outcome(
5596+
&self,
5597+
session_id: String,
5598+
) -> BitFunResult<ContextCompactionOutcome> {
55875599
let task = self.start_manual_compaction_task(session_id, None).await?;
55885600
task.completion.await.map_err(|_| {
55895601
BitFunError::Service(format!(

src/crates/assembly/core/src/agentic/tools/implementations/session_control_tool.rs

Lines changed: 172 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ impl SessionControlTool {
104104
match action {
105105
SessionControlAction::Cancel
106106
| SessionControlAction::Delete
107+
| SessionControlAction::Compact
107108
| SessionControlAction::Rename => {
108109
let session_id = session_id.ok_or_else(|| {
109110
BitFunError::tool(format!("session_id is required for {}", action.as_str()))
@@ -274,6 +275,7 @@ Actions:
274275
- "create": Create a new session. You may optionally provide session_name and agent_type.
275276
- "cancel": Cancel the target session's currently running dialog turn. This does not delete the session or clear any queued messages that may still run later.
276277
- "delete": Delete an existing session by session_id.
278+
- "compact": Compact the target session's context to reduce its token footprint. Returns the applied compaction statistics.
277279
- "rename": Rename an existing session by session_id using session_name as the new title.
278280
- "list": List all sessions.
279281
@@ -285,13 +287,13 @@ Arguments:
285287
- "Plan": Planning agent for clarifying requirements and producing an implementation plan before coding.
286288
- "Cowork": Collaborative agent for office-style work such as research, documentation, presentations, etc.
287289
- "DeepResearch": Research agent for systematic investigation and evidence-driven reports.
288-
- "session_id": Required for cancel, delete, and rename."#
290+
- "session_id": Required for cancel, delete, compact, and rename."#
289291
.to_string(),
290292
)
291293
}
292294

293295
fn short_description(&self) -> String {
294-
"Create, list, rename, cancel, and delete persisted agent sessions.".to_string()
296+
"Create, list, rename, compact, cancel, and delete persisted agent sessions.".to_string()
295297
}
296298

297299
fn default_exposure(&self) -> ToolExposure {
@@ -304,7 +306,7 @@ Arguments:
304306
"properties": {
305307
"action": {
306308
"type": "string",
307-
"enum": ["create", "cancel", "delete", "rename", "list"],
309+
"enum": ["create", "cancel", "delete", "compact", "rename", "list"],
308310
"description": "The session action to perform."
309311
},
310312
"workspace": {
@@ -313,7 +315,7 @@ Arguments:
313315
},
314316
"session_id": {
315317
"type": "string",
316-
"description": "Required for cancel, delete, and rename."
318+
"description": "Required for cancel, delete, compact, and rename."
317319
},
318320
"session_name": {
319321
"type": "string",
@@ -642,6 +644,113 @@ Arguments:
642644
image_attachments: None,
643645
}])
644646
}
647+
SessionControlAction::Compact => {
648+
let session_id = params.session_id.as_deref().ok_or_else(|| {
649+
BitFunError::tool("session_id is required for compact".to_string())
650+
})?;
651+
validate_session_id(session_id).map_err(BitFunError::tool)?;
652+
let workspace = self
653+
.resolve_effective_workspace(
654+
SessionControlAction::Compact,
655+
Some(session_id),
656+
context,
657+
&runtime,
658+
)
659+
.await?;
660+
661+
// Authorization follows the owner/creator/ancestor semantics and
662+
// additionally permits compacting the caller's own session.
663+
let current_session_id = context.session_id.as_ref().ok_or_else(|| {
664+
BitFunError::tool(
665+
"cannot compact a session without a caller session in tool context"
666+
.to_string(),
667+
)
668+
})?;
669+
let session_manager = coordinator.get_session_manager();
670+
let caller_is_owner = session_manager
671+
.get_session(current_session_id)
672+
.is_some_and(|session| session.created_by.is_none());
673+
let is_self = current_session_id == session_id;
674+
let created_by_match = session_manager
675+
.load_session_metadata(
676+
std::path::Path::new(&workspace.project_workspace),
677+
session_id,
678+
)
679+
.await
680+
.ok()
681+
.flatten()
682+
.and_then(|metadata| metadata.created_by)
683+
.is_some_and(|creator| {
684+
creator == session_control_creator_marker(current_session_id)
685+
});
686+
if !caller_is_owner && !is_self && !created_by_match {
687+
let mut ancestors = Vec::new();
688+
let mut visited = std::collections::HashSet::new();
689+
visited.insert(session_id.to_string());
690+
let mut current = session_id.to_string();
691+
loop {
692+
let metadata = session_manager
693+
.load_session_metadata(
694+
std::path::Path::new(&workspace.project_workspace),
695+
&current,
696+
)
697+
.await
698+
.ok()
699+
.flatten();
700+
match metadata
701+
.and_then(|m| m.relationship.and_then(|r| r.parent_session_id))
702+
{
703+
Some(parent_id) => {
704+
if !visited.insert(parent_id.clone()) {
705+
break;
706+
}
707+
ancestors.push(parent_id.clone());
708+
current = parent_id;
709+
}
710+
None => break,
711+
}
712+
}
713+
if !ancestors.is_empty() && !ancestors.contains(current_session_id) {
714+
return Err(BitFunError::tool(format!(
715+
"session '{current_session_id}' is not authorized to compact session '{session_id}': not a parent/ancestor and not the creator"
716+
)));
717+
}
718+
}
719+
720+
let outcome = coordinator
721+
.compact_session_with_outcome(session_id.to_string())
722+
.await
723+
.map_err(|error| {
724+
BitFunError::tool(format!(
725+
"cannot compact session '{session_id}': {}",
726+
error
727+
))
728+
})?;
729+
730+
Ok(vec![ToolResult::Result {
731+
data: json!({
732+
"success": true,
733+
"action": "compact",
734+
"workspace": workspace.display_workspace.clone(),
735+
"session_id": session_id,
736+
"applied": outcome.applied,
737+
"tokens_before": outcome.tokens_before,
738+
"tokens_after": outcome.tokens_after,
739+
"compression_ratio": outcome.compression_ratio,
740+
"duration": outcome.duration_ms,
741+
"summary_source": if outcome.has_summary {
742+
Some(outcome.summary_source)
743+
} else {
744+
None
745+
},
746+
}),
747+
result_for_assistant: Some(format!(
748+
"Compacted session '{session_id}' in workspace '{}'.",
749+
workspace.display_workspace
750+
)),
751+
image_attachments: None,
752+
}])
753+
}
645754
SessionControlAction::List => {
646755
let workspace = self
647756
.resolve_effective_workspace(
@@ -956,4 +1065,63 @@ mod tests {
9561065

9571066
assert!(validation.result, "{:?}", validation.message);
9581067
}
1068+
1069+
#[tokio::test]
1070+
async fn validate_compact_requires_session_id() {
1071+
let tool = SessionControlTool::new();
1072+
1073+
let validation = tool
1074+
.validate_input(
1075+
&json!({
1076+
"action": "compact",
1077+
}),
1078+
Some(&empty_context()),
1079+
)
1080+
.await;
1081+
1082+
assert!(!validation.result);
1083+
assert_eq!(
1084+
validation.message.as_deref(),
1085+
Some("session_id is required for compact")
1086+
);
1087+
}
1088+
1089+
#[tokio::test]
1090+
async fn validate_compact_rejects_session_name() {
1091+
let tool = SessionControlTool::new();
1092+
1093+
let validation = tool
1094+
.validate_input(
1095+
&json!({
1096+
"action": "compact",
1097+
"session_id": "worker_1",
1098+
"session_name": "should-not-be-here",
1099+
}),
1100+
Some(&empty_context()),
1101+
)
1102+
.await;
1103+
1104+
assert!(!validation.result);
1105+
assert_eq!(
1106+
validation.message.as_deref(),
1107+
Some("session_name is only allowed for create")
1108+
);
1109+
}
1110+
1111+
#[tokio::test]
1112+
async fn validate_compact_accepts_session_id() {
1113+
let tool = SessionControlTool::new();
1114+
1115+
let validation = tool
1116+
.validate_input(
1117+
&json!({
1118+
"action": "compact",
1119+
"session_id": "worker_1",
1120+
}),
1121+
Some(&empty_context()),
1122+
)
1123+
.await;
1124+
1125+
assert!(validation.result, "{:?}", validation.message);
1126+
}
9591127
}

src/crates/execution/agent-runtime/src/session_control.rs

Lines changed: 85 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ pub enum SessionControlAction {
1111
Cancel,
1212
Delete,
1313
List,
14+
Compact,
1415
Rename,
1516
}
1617

@@ -21,6 +22,7 @@ impl SessionControlAction {
2122
Self::Cancel => "cancel",
2223
Self::Delete => "delete",
2324
Self::List => "list",
25+
Self::Compact => "compact",
2426
Self::Rename => "rename",
2527
}
2628
}
@@ -183,7 +185,13 @@ fn validate_mutating_action_target(
183185
return invalid(message);
184186
}
185187

186-
if context.current_session_id == Some(session_id) && context.has_workspace_root {
188+
// Guard only depends on session-binding equivalence: if the target is the
189+
// current session it is refused. `compact` is the sole exception (it may
190+
// compress the current session and resident subagent workstations).
191+
if !matches!(action, SessionControlAction::Compact)
192+
&& context.current_session_id == Some(session_id)
193+
&& context.has_workspace_root
194+
{
187195
return invalid(format!(
188196
"cannot {} the current session from SessionControl",
189197
action.as_str()
@@ -226,6 +234,7 @@ pub fn validate_session_control_input(
226234
}
227235
SessionControlAction::Cancel
228236
| SessionControlAction::Delete
237+
| SessionControlAction::Compact
229238
| SessionControlAction::Rename => {
230239
return validate_mutating_action_target(&input.action, input, context);
231240
}
@@ -266,6 +275,7 @@ pub fn render_session_control_tool_use_message(input: &Value) -> String {
266275
"create" => format!("Create session in {workspace}"),
267276
"cancel" => format!("Cancel active turn for session {session_id}"),
268277
"delete" => format!("Delete session {session_id}"),
278+
"compact" => format!("Compact session {session_id}"),
269279
"rename" => format!("Rename session {session_id}"),
270280
"list" => format!("List sessions in {workspace}"),
271281
_ => format!("Manage sessions in {workspace}"),
@@ -443,4 +453,78 @@ mod tests {
443453
assert!(message.contains("new-title"));
444454
assert!(message.contains("/ws"));
445455
}
456+
457+
#[test]
458+
fn compact_action_parses_payload_session_id() {
459+
let input: SessionControlInput = serde_json::from_value(json!({
460+
"action": "compact",
461+
"session_id": "worker_1",
462+
}))
463+
.expect("compact payload must parse");
464+
assert_eq!(input.action, SessionControlAction::Compact);
465+
assert_eq!(input.session_id.as_deref(), Some("worker_1"));
466+
assert_eq!(SessionControlAction::Compact.as_str(), "compact");
467+
}
468+
469+
#[test]
470+
fn compact_validation_requires_session_id() {
471+
let input = SessionControlInput {
472+
action: SessionControlAction::Compact,
473+
workspace: None,
474+
session_id: None,
475+
session_name: None,
476+
agent_type: None,
477+
};
478+
let result = validate_session_control_input(&input, context(None));
479+
assert!(!result.result);
480+
assert!(result
481+
.message
482+
.as_deref()
483+
.unwrap_or_default()
484+
.contains("session_id is required"));
485+
}
486+
487+
#[test]
488+
fn compact_validation_rejects_non_mutating_fields() {
489+
let input = SessionControlInput {
490+
action: SessionControlAction::Compact,
491+
workspace: None,
492+
session_id: Some("worker_1".to_string()),
493+
session_name: Some("should not be allowed".to_string()),
494+
agent_type: None,
495+
};
496+
let result = validate_session_control_input(&input, context(None));
497+
assert!(!result.result);
498+
assert_eq!(
499+
result.message.as_deref(),
500+
Some("session_name is only allowed for create")
501+
);
502+
}
503+
504+
#[test]
505+
fn compact_validation_allows_current_session() {
506+
let input = SessionControlInput {
507+
action: SessionControlAction::Compact,
508+
workspace: None,
509+
session_id: Some("self_1".to_string()),
510+
session_name: None,
511+
agent_type: None,
512+
};
513+
let result = validate_session_control_input(&input, context(Some("self_1")));
514+
assert!(
515+
result.result,
516+
"compact of the current session must be allowed: {:?}",
517+
result.message
518+
);
519+
}
520+
521+
#[test]
522+
fn compact_render_mentions_session() {
523+
let rendered = render_session_control_tool_use_message(&json!({
524+
"action": "compact",
525+
"session_id": "worker_1",
526+
}));
527+
assert!(rendered.contains("Compact session"));
528+
assert!(rendered.contains("worker_1"));
529+
}
446530
}

0 commit comments

Comments
 (0)