1use super::{Capability, CapabilityLocalization, CapabilityStatus};
8use crate::error::{AgentLoopError, Result};
9use crate::events::{EventContext, EventRequest, SessionTitleUpdatedData, TokenUsage};
10use crate::session::Session;
11use crate::tool_types::ToolHints;
12use crate::tools::{Tool, ToolExecutionResult};
13use crate::traits::{EventEmitter, SessionMutator, SessionStore, ToolContext};
14use crate::typed_id::SessionId;
15use async_trait::async_trait;
16use serde::{Deserialize, Serialize};
17use serde_json::{Value, json};
18
19pub const SESSION_CAPABILITY_ID: &str = "session";
20
21#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
23#[serde(deny_unknown_fields)]
24pub struct SessionCapabilityConfig {
25 #[serde(default)]
27 pub auto_title: bool,
28}
29
30impl SessionCapabilityConfig {
31 fn from_value(config: &Value) -> Self {
32 serde_json::from_value(config.clone()).unwrap_or_default()
33 }
34}
35
36#[derive(Debug, Clone)]
38pub struct SessionTitleMutation {
39 pub session: Session,
41 pub changed: bool,
43}
44
45pub fn session_title_updated_event(
50 session_id: SessionId,
51 event_context: EventContext,
52 previous_title: Option<String>,
53 title: String,
54) -> Option<EventRequest> {
55 if previous_title.as_deref() == Some(title.as_str()) {
56 return None;
57 }
58 Some(EventRequest::new(
59 session_id,
60 event_context,
61 SessionTitleUpdatedData {
62 previous_title,
63 title,
64 },
65 ))
66}
67
68pub async fn update_session_title_with_event(
75 session_id: SessionId,
76 title: String,
77 event_context: EventContext,
78 session_store: &dyn SessionStore,
79 session_mutator: &dyn SessionMutator,
80 event_emitter: &dyn EventEmitter,
81) -> Result<SessionTitleMutation> {
82 let current = session_store
83 .get_session(session_id)
84 .await?
85 .ok_or_else(|| AgentLoopError::store(format!("session not found: {session_id}")))?;
86 let previous_title = current.title.clone();
87
88 let Some(event) =
89 session_title_updated_event(session_id, event_context, previous_title, title.clone())
90 else {
91 return Ok(SessionTitleMutation {
92 session: current,
93 changed: false,
94 });
95 };
96
97 let session = session_mutator
98 .update_session_title(session_id, title.clone())
99 .await?;
100 event_emitter.emit(event).await?;
101
102 Ok(SessionTitleMutation {
103 session,
104 changed: true,
105 })
106}
107
108pub struct SessionCapability;
110
111#[async_trait]
112impl Capability for SessionCapability {
113 fn id(&self) -> &str {
114 SESSION_CAPABILITY_ID
115 }
116
117 fn name(&self) -> &str {
118 "Session"
119 }
120
121 fn description(&self) -> &str {
122 "Read and update current session metadata like title and agent info."
123 }
124
125 fn localizations(&self) -> Vec<CapabilityLocalization> {
126 vec![CapabilityLocalization::text(
127 "uk",
128 "Сесія",
129 "Читання та оновлення метаданих поточної сесії, як-от назви й інформації про агента.",
130 )]
131 }
132
133 fn status(&self) -> CapabilityStatus {
134 CapabilityStatus::Available
135 }
136
137 fn icon(&self) -> Option<&str> {
138 Some("panel-left")
139 }
140
141 fn category(&self) -> Option<&str> {
142 Some("Session")
143 }
144
145 fn config_schema(&self) -> Option<Value> {
146 Some(json!({
147 "type": "object",
148 "properties": {
149 "auto_title": {
150 "type": "boolean",
151 "title": "Automatic session titles",
152 "description": "Require a concise title before handling the first substantive request and update it only when the conversation's primary theme materially changes.",
153 "default": false
154 }
155 },
156 "additionalProperties": false
157 }))
158 }
159
160 fn validate_config(&self, config: &Value) -> std::result::Result<(), String> {
161 if config.is_null() {
162 return Ok(());
163 }
164 serde_json::from_value::<SessionCapabilityConfig>(config.clone())
165 .map(|_| ())
166 .map_err(|error| format!("invalid session config: {error}"))
167 }
168
169 async fn system_prompt_contribution_with_config(
170 &self,
171 _ctx: &super::SystemPromptContext,
172 config: &Value,
173 ) -> Option<String> {
174 if !SessionCapabilityConfig::from_value(config).auto_title {
175 return None;
176 }
177
178 Some(format!(
179 "<capability id=\"{}\">\nTitle maintenance is mandatory when automatic titles are enabled. On the first substantive user request, you MUST call `write_session_title` with a concise 3–7 word title for the conversation's primary theme before using any other tool, doing substantive work, or giving a substantive response. Ignore greetings, acknowledgements, and other non-substantive messages. On later turns, if the primary theme materially changes, you MUST call `write_session_title` before using any other tool, doing substantive work, or responding. You MUST NOT update the title for minor subtopics, follow-ups, or implementation details. Calling `write_session_title` updates session metadata only; it does not change project or workspace files and must not be treated as a project-file change.\n</capability>",
180 self.id()
181 ))
182 }
183
184 fn tools(&self) -> Vec<Box<dyn Tool>> {
185 vec![
186 Box::new(WriteSessionTitleTool),
187 Box::new(GetSessionInfoTool),
188 ]
189 }
190}
191
192pub struct WriteSessionTitleTool;
194
195#[async_trait]
196impl Tool for WriteSessionTitleTool {
197 fn narrate(
198 &self,
199 tool_call: &crate::tool_types::ToolCall,
200 phase: crate::tool_narration::ToolNarrationPhase,
201 locale: Option<&str>,
202 _ctx: crate::tool_narration::ToolNarrationContext<'_>,
203 ) -> Option<String> {
204 Some(crate::tool_narration::narrate_write_session_title(
205 &tool_call.arguments,
206 phase,
207 locale,
208 ))
209 }
210
211 fn name(&self) -> &str {
212 "write_session_title"
213 }
214
215 fn display_name(&self) -> Option<&str> {
216 Some("Write Session Title")
217 }
218
219 fn description(&self) -> &str {
220 "Update the current session title."
221 }
222
223 fn parameters_schema(&self) -> Value {
224 json!({
225 "type": "object",
226 "properties": {
227 "title": {
228 "type": "string",
229 "description": "New session title"
230 }
231 },
232 "required": ["title"],
233 "additionalProperties": false
234 })
235 }
236
237 fn hints(&self) -> ToolHints {
238 ToolHints::default().with_idempotent(true)
239 }
240
241 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
242 ToolExecutionResult::tool_error(
243 "write_session_title requires context. This tool must be executed with session context.",
244 )
245 }
246
247 async fn execute_with_context(
248 &self,
249 arguments: Value,
250 context: &ToolContext,
251 ) -> ToolExecutionResult {
252 let title = match arguments.get("title").and_then(|v| v.as_str()) {
253 Some(t) if !t.trim().is_empty() => t.trim().to_string(),
254 _ => return ToolExecutionResult::tool_error("Missing required parameter: title"),
255 };
256
257 let Some(session_store) = &context.session_store else {
258 return ToolExecutionResult::tool_error("Session store not available in this context");
259 };
260 let Some(mutator) = &context.session_mutator else {
261 return ToolExecutionResult::tool_error(
262 "Session mutator not available in this context",
263 );
264 };
265 let Some(event_emitter) = &context.event_emitter else {
266 return ToolExecutionResult::tool_error("Event emitter not available in this context");
267 };
268
269 match update_session_title_with_event(
270 context.session_id,
271 title,
272 context.event_context.clone().unwrap_or_default(),
273 session_store.as_ref(),
274 mutator.as_ref(),
275 event_emitter.as_ref(),
276 )
277 .await
278 {
279 Ok(outcome) => ToolExecutionResult::success(json!({
280 "session_id": outcome.session.id.to_string(),
281 "title": outcome.session.title,
282 "updated": outcome.changed,
283 })),
284 Err(e) => ToolExecutionResult::internal_error(e),
285 }
286 }
287}
288
289pub struct GetSessionInfoTool;
291
292#[async_trait]
293impl Tool for GetSessionInfoTool {
294 fn narrate(
295 &self,
296 _tool_call: &crate::tool_types::ToolCall,
297 phase: crate::tool_narration::ToolNarrationPhase,
298 locale: Option<&str>,
299 _ctx: crate::tool_narration::ToolNarrationContext<'_>,
300 ) -> Option<String> {
301 Some(crate::tool_narration::narrate_get_session_info(
302 phase, locale,
303 ))
304 }
305
306 fn name(&self) -> &str {
307 "get_session_info"
308 }
309
310 fn display_name(&self) -> Option<&str> {
311 Some("Get Session Info")
312 }
313
314 fn description(&self) -> &str {
315 "Get current session metadata: id, title, locale, agent name, and cumulative token usage."
316 }
317
318 fn parameters_schema(&self) -> Value {
319 json!({
320 "type": "object",
321 "properties": {},
322 "additionalProperties": false
323 })
324 }
325
326 fn hints(&self) -> ToolHints {
327 ToolHints::default()
328 .with_readonly(true)
329 .with_idempotent(true)
330 }
331
332 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
333 ToolExecutionResult::tool_error(
334 "get_session_info requires context. This tool must be executed with session context.",
335 )
336 }
337
338 async fn execute_with_context(
339 &self,
340 _arguments: Value,
341 context: &ToolContext,
342 ) -> ToolExecutionResult {
343 let Some(session_store) = &context.session_store else {
344 return ToolExecutionResult::tool_error("Session store not available in this context");
345 };
346
347 let session = match session_store.get_session(context.session_id).await {
348 Ok(Some(session)) => session,
349 Ok(None) => return ToolExecutionResult::tool_error("Session not found"),
350 Err(e) => return ToolExecutionResult::internal_error(e),
351 };
352
353 let agent_name = if let (Some(agent_id), Some(agent_store)) =
354 (session.agent_id, &context.agent_store)
355 {
356 match agent_store.get_agent(agent_id).await {
357 Ok(Some(agent)) => Some(agent.display_name.unwrap_or_else(|| agent.name.clone())),
358 Ok(None) => None,
359 Err(e) => return ToolExecutionResult::internal_error(e),
360 }
361 } else {
362 None
363 };
364
365 ToolExecutionResult::success(json!({
366 "session_id": session.id.to_string(),
367 "title": session.title,
368 "locale": session.locale,
369 "agent_name": agent_name,
370 "usage": session.usage.as_ref().map(usage_json),
371 }))
372 }
373}
374
375fn usage_json(usage: &TokenUsage) -> Value {
376 json!({
377 "input_tokens": usage.input_tokens,
378 "output_tokens": usage.output_tokens,
379 "cache_read_tokens": usage.cache_read_tokens,
380 "cache_creation_tokens": usage.cache_creation_tokens,
381 "total_tokens": usage.total_tokens(),
382 })
383}
384
385#[cfg(test)]
386mod tests {
387 use super::*;
388 use crate::agent::{Agent, AgentStatus};
389 use crate::events::{Event, EventRequest};
390 use crate::session::{Session, SessionStatus};
391 use crate::typed_id::{AgentId, EventId, HarnessId, MessageId, ModelId, SessionId, TurnId};
392 use crate::{AgentCapabilityConfig, Tool};
393 use async_trait::async_trait;
394 use chrono::Utc;
395 use std::sync::{Arc, Mutex};
396
397 #[derive(Clone)]
398 struct MockSessionStore {
399 session: Arc<Mutex<Option<Session>>>,
400 }
401
402 #[async_trait]
403 impl crate::traits::SessionStore for MockSessionStore {
404 async fn get_session(&self, _session_id: SessionId) -> Result<Option<Session>> {
405 Ok(self.session.lock().expect("poisoned").clone())
406 }
407 }
408
409 #[derive(Clone)]
410 struct MockSessionMutator {
411 session: Arc<Mutex<Session>>,
412 }
413
414 #[async_trait]
415 impl crate::traits::SessionMutator for MockSessionMutator {
416 async fn update_session_title(
417 &self,
418 _session_id: SessionId,
419 title: String,
420 ) -> Result<Session> {
421 let mut session = self.session.lock().expect("poisoned");
422 session.title = Some(title);
423 Ok(session.clone())
424 }
425 }
426
427 struct MockAgentStore {
428 agent: Option<Agent>,
429 }
430
431 #[derive(Clone, Default)]
432 struct RecordingEventEmitter {
433 requests: Arc<Mutex<Vec<EventRequest>>>,
434 }
435
436 #[async_trait]
437 impl crate::traits::EventEmitter for RecordingEventEmitter {
438 async fn emit(&self, request: EventRequest) -> Result<Event> {
439 self.requests
440 .lock()
441 .expect("poisoned")
442 .push(request.clone());
443 Ok(request.into_event(EventId::new(), 1))
444 }
445 }
446
447 #[async_trait]
448 impl crate::traits::AgentStore for MockAgentStore {
449 async fn get_agent(&self, _agent_id: AgentId) -> Result<Option<Agent>> {
450 Ok(self.agent.clone())
451 }
452 }
453
454 fn build_session(agent_id: Option<AgentId>) -> Session {
455 let session_id = SessionId::new();
456 Session {
457 id: session_id,
458 workspace_id: crate::WorkspaceId::from_uuid(session_id.uuid()),
460 organization_id: "org_00000000000000000000000000000001".to_string(),
461 harness_id: HarnessId::new(),
462 agent_id,
463 agent_version_id: None,
464 agent_identity_id: None,
465 owner_principal_id: crate::PrincipalId::from_seed(1),
466 resolved_owner_user_id: None,
467 owner: None,
468 effective_owner: None,
469 title: Some("Old title".to_string()),
470 goal: None,
471 locale: None,
472 preview: None,
473 output_preview: None,
474 tags: vec![],
475 model_id: Some(ModelId::new()),
476 capabilities: vec![],
477 tools: vec![],
478 mcp_servers: Default::default(),
479 system_prompt: None,
480 initial_files: vec![],
481 hints: None,
482 network_access: None,
483 max_iterations: None,
484 parallel_tool_calls: None,
485 status: SessionStatus::Idle,
486 created_at: Utc::now(),
487 updated_at: Utc::now(),
488 started_at: None,
489 finished_at: None,
490 usage: None,
491 is_pinned: None,
492 active_schedule_count: None,
493 features: vec![],
494 parent_session_id: None,
495 forked_from_session_id: None,
496 forked_from_sequence: None,
497 blueprint_id: None,
498 blueprint_config: None,
499 }
500 }
501
502 #[test]
503 fn session_tools_narrate_all_phases() {
504 use crate::tool_narration::{ToolNarrationContext, ToolNarrationPhase};
505 use crate::tool_types::ToolCall;
506
507 let ctx = ToolNarrationContext::default();
508 let title_call = ToolCall {
509 id: "c1".to_string(),
510 name: "write_session_title".to_string(),
511 arguments: json!({ "title": "Improve tool narration" }),
512 };
513 assert_eq!(
514 WriteSessionTitleTool.narrate(&title_call, ToolNarrationPhase::Started, None, ctx),
515 Some("Updating session title: Improve tool narration".to_string())
516 );
517 assert_eq!(
518 WriteSessionTitleTool.narrate(&title_call, ToolNarrationPhase::Completed, None, ctx),
519 Some("Updated session title: Improve tool narration".to_string())
520 );
521 assert_eq!(
522 WriteSessionTitleTool.narrate(&title_call, ToolNarrationPhase::Failed, None, ctx),
523 Some("Failed to update session title: Improve tool narration".to_string())
524 );
525
526 let cap = SessionCapability;
528 let def = WriteSessionTitleTool.to_definition();
529 assert_eq!(
530 cap.narrate(
531 Some(&def),
532 &title_call,
533 ToolNarrationPhase::Started,
534 None,
535 ctx
536 ),
537 Some("Updating session title: Improve tool narration".to_string())
538 );
539
540 let info_call = ToolCall {
541 id: "c2".to_string(),
542 name: "get_session_info".to_string(),
543 arguments: json!({}),
544 };
545 assert_eq!(
546 GetSessionInfoTool.narrate(&info_call, ToolNarrationPhase::Completed, None, ctx),
547 Some("Read session info".to_string())
548 );
549 }
550
551 #[tokio::test]
552 async fn write_session_title_updates_title() {
553 let session = build_session(None);
554 let session_id = session.id;
555 let stored = Arc::new(Mutex::new(Some(session.clone())));
556 let emitter = RecordingEventEmitter::default();
557 let turn_id = TurnId::new();
558 let input_message_id = MessageId::new();
559 let mut context = ToolContext::new(session_id);
560 context.session_store = Some(Arc::new(MockSessionStore { session: stored }));
561 context.session_mutator = Some(Arc::new(MockSessionMutator {
562 session: Arc::new(Mutex::new(session)),
563 }));
564 context.event_emitter = Some(Arc::new(emitter.clone()));
565 context.event_context = Some(EventContext::turn(turn_id, input_message_id));
566
567 let tool = WriteSessionTitleTool;
568 let result = tool
569 .execute_with_context(json!({"title": "New title"}), &context)
570 .await;
571
572 match result {
573 ToolExecutionResult::Success(value) => {
574 assert_eq!(value["title"], "New title");
575 assert_eq!(value["updated"], true);
576 }
577 _ => panic!("expected success"),
578 }
579
580 let requests = emitter.requests.lock().expect("poisoned");
581 assert_eq!(requests.len(), 1);
582 assert_eq!(requests[0].event_type, crate::events::SESSION_TITLE_UPDATED);
583 assert_eq!(requests[0].context.turn_id, Some(turn_id));
584 assert_eq!(requests[0].context.input_message_id, Some(input_message_id));
585 match &requests[0].data {
586 crate::events::EventData::SessionTitleUpdated(data) => {
587 assert_eq!(data.previous_title.as_deref(), Some("Old title"));
588 assert_eq!(data.title, "New title");
589 }
590 data => panic!("unexpected event data: {data:?}"),
591 }
592 }
593
594 #[tokio::test]
595 async fn write_session_title_is_noop_when_title_is_unchanged() {
596 let session = build_session(None);
597 let session_id = session.id;
598 let emitter = RecordingEventEmitter::default();
599 let mut context = ToolContext::new(session_id);
600 context.session_store = Some(Arc::new(MockSessionStore {
601 session: Arc::new(Mutex::new(Some(session.clone()))),
602 }));
603 context.session_mutator = Some(Arc::new(MockSessionMutator {
604 session: Arc::new(Mutex::new(session)),
605 }));
606 context.event_emitter = Some(Arc::new(emitter.clone()));
607
608 let result = WriteSessionTitleTool
609 .execute_with_context(json!({"title": "Old title"}), &context)
610 .await;
611
612 match result {
613 ToolExecutionResult::Success(value) => assert_eq!(value["updated"], false),
614 _ => panic!("expected success"),
615 }
616 assert!(emitter.requests.lock().expect("poisoned").is_empty());
617 }
618
619 #[tokio::test]
620 async fn auto_title_policy_is_opt_in_and_mandatory_when_enabled() {
621 let capability = SessionCapability;
622 let ctx = super::super::SystemPromptContext::without_file_store(SessionId::new());
623
624 assert!(
625 capability
626 .system_prompt_contribution_with_config(&ctx, &json!({}))
627 .await
628 .is_none()
629 );
630 let prompt = capability
631 .system_prompt_contribution_with_config(&ctx, &json!({"auto_title": true}))
632 .await
633 .expect("auto-title prompt");
634 assert_eq!(
635 prompt,
636 "<capability id=\"session\">\nTitle maintenance is mandatory when automatic titles are enabled. On the first substantive user request, you MUST call `write_session_title` with a concise 3–7 word title for the conversation's primary theme before using any other tool, doing substantive work, or giving a substantive response. Ignore greetings, acknowledgements, and other non-substantive messages. On later turns, if the primary theme materially changes, you MUST call `write_session_title` before using any other tool, doing substantive work, or responding. You MUST NOT update the title for minor subtopics, follow-ups, or implementation details. Calling `write_session_title` updates session metadata only; it does not change project or workspace files and must not be treated as a project-file change.\n</capability>"
637 );
638 }
639
640 #[tokio::test]
641 async fn get_session_info_returns_agent_name_when_assigned() {
642 let agent_id = AgentId::new();
643 let session = build_session(Some(agent_id));
644 let session_id = session.id;
645
646 let agent = Agent {
647 public_id: agent_id,
648 internal_id: agent_id.uuid(),
649 name: "research-agent".to_string(),
650 display_name: Some("Research Agent".to_string()),
651 description: Some("desc".to_string()),
652 system_prompt: "prompt".to_string(),
653 default_model_id: None,
654
655 harness_id: crate::typed_id::HarnessId::from_uuid(uuid::Uuid::nil()),
656 default_version_id: None,
657 forked_from_agent_id: None,
658 forked_from_version_id: None,
659 root_agent_id: None,
660 tags: vec![],
661 capabilities: vec![AgentCapabilityConfig::new("session")],
662 initial_files: vec![],
663 network_access: None,
664 max_iterations: None,
665 parallel_tool_calls: None,
666 tools: vec![],
667 mcp_servers: Default::default(),
668 status: AgentStatus::Active,
669 created_at: Utc::now(),
670 updated_at: Utc::now(),
671 archived_at: None,
672 deleted_at: None,
673 usage: None,
674 };
675
676 let context = ToolContext::new(session_id)
677 .with_session_store(Arc::new(MockSessionStore {
678 session: Arc::new(Mutex::new(Some(session))),
679 }))
680 .with_agent_store(Arc::new(MockAgentStore { agent: Some(agent) }));
681
682 let tool = GetSessionInfoTool;
683 let result = tool.execute_with_context(json!({}), &context).await;
684
685 match result {
686 ToolExecutionResult::Success(value) => {
687 assert_eq!(value["title"], "Old title");
688 assert_eq!(value["agent_name"], "Research Agent");
689 assert!(value["usage"].is_null());
690 }
691 _ => panic!("expected success"),
692 }
693 }
694
695 #[tokio::test]
696 async fn get_session_info_returns_cumulative_usage() {
697 let mut session = build_session(None);
698 session.usage = Some(TokenUsage::with_cache(120, 45, Some(30), Some(10)));
699 let session_id = session.id;
700
701 let context = ToolContext::new(session_id).with_session_store(Arc::new(MockSessionStore {
702 session: Arc::new(Mutex::new(Some(session))),
703 }));
704
705 let tool = GetSessionInfoTool;
706 let result = tool.execute_with_context(json!({}), &context).await;
707
708 match result {
709 ToolExecutionResult::Success(value) => {
710 assert_eq!(value["usage"]["input_tokens"], 120);
711 assert_eq!(value["usage"]["output_tokens"], 45);
712 assert_eq!(value["usage"]["cache_read_tokens"], 30);
713 assert_eq!(value["usage"]["cache_creation_tokens"], 10);
714 assert_eq!(value["usage"]["total_tokens"], 165);
715 }
716 _ => panic!("expected success"),
717 }
718 }
719}