use adk_core::{
Agent, AppName, CallbackContext, Content, Event, FunctionResponseData,
InvocationContext as InvocationContextTrait, Part, ReadonlyContext, RunConfig,
Session as CoreSession, SessionId, StreamingMode, UserId,
};
use adk_runner::{InvocationContext, MutableSession};
use adk_session::{Events, Session, State};
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
struct MockEvents;
impl Events for MockEvents {
fn all(&self) -> Vec<Event> {
Vec::new()
}
fn len(&self) -> usize {
0
}
fn at(&self, _index: usize) -> Option<&Event> {
None
}
}
struct MockSessionStateView(std::sync::Arc<std::sync::RwLock<HashMap<String, serde_json::Value>>>);
impl State for MockSessionStateView {
fn get(&self, key: &str) -> Option<serde_json::Value> {
self.0.read().unwrap().get(key).cloned()
}
fn set(&mut self, key: String, value: serde_json::Value) {
self.0.write().unwrap().insert(key, value);
}
fn all(&self) -> HashMap<String, serde_json::Value> {
self.0.read().unwrap().clone()
}
}
struct MockSessionWithState {
state_view: MockSessionStateView,
}
impl MockSessionWithState {
fn new() -> Self {
let state_arc = std::sync::Arc::new(std::sync::RwLock::new(HashMap::new()));
Self { state_view: MockSessionStateView(state_arc) }
}
fn with_state(state: HashMap<String, serde_json::Value>) -> Self {
let state_arc = std::sync::Arc::new(std::sync::RwLock::new(state));
Self { state_view: MockSessionStateView(state_arc) }
}
}
impl Session for MockSessionWithState {
fn id(&self) -> &str {
"session-789"
}
fn app_name(&self) -> &str {
"test-app"
}
fn user_id(&self) -> &str {
"user-456"
}
fn state(&self) -> &dyn State {
&self.state_view
}
fn events(&self) -> &dyn adk_session::Events {
static MOCK_EVENTS: MockEvents = MockEvents;
&MOCK_EVENTS
}
fn last_update_time(&self) -> chrono::DateTime<chrono::Utc> {
chrono::Utc::now()
}
}
struct MockAgent {
name: String,
}
#[async_trait]
impl Agent for MockAgent {
fn name(&self) -> &str {
&self.name
}
fn description(&self) -> &str {
"Mock agent"
}
fn sub_agents(&self) -> &[Arc<dyn Agent>] {
&[]
}
async fn run(
&self,
_ctx: Arc<dyn InvocationContextTrait>,
) -> adk_core::Result<adk_core::EventStream> {
Ok(Box::pin(futures::stream::empty()))
}
}
#[test]
fn test_context_creation() {
let agent = Arc::new(MockAgent { name: "test_agent".to_string() });
let content =
Content { role: "user".to_string(), parts: vec![Part::Text { text: "Hello".to_string() }] };
let ctx = InvocationContext::new(
"inv-123".to_string(),
agent.clone(),
"user-456".to_string(),
"test-app".to_string(),
"session-789".to_string(),
content.clone(),
Arc::new(MockSessionWithState::new()),
)
.unwrap();
assert_eq!(ctx.invocation_id(), "inv-123");
assert_eq!(ctx.agent_name(), "test_agent");
assert_eq!(ctx.user_id(), "user-456");
assert_eq!(ctx.app_name(), "test-app");
assert_eq!(ctx.session_id(), "session-789");
assert_eq!(ctx.branch(), "");
assert_eq!(ctx.user_content().role, "user");
}
#[test]
fn test_context_creation_with_typed_ids() {
let agent = Arc::new(MockAgent { name: "test_agent".to_string() });
let content = Content::new("user");
let ctx = InvocationContext::new_typed(
"inv-typed".to_string(),
agent,
UserId::new("user-456").unwrap(),
AppName::new("test-app").unwrap(),
SessionId::new("session-789").unwrap(),
content,
Arc::new(MockSessionWithState::new()),
)
.unwrap();
assert_eq!(ctx.invocation_id(), "inv-typed");
assert_eq!(ctx.user_id(), "user-456");
assert_eq!(ctx.app_name(), "test-app");
assert_eq!(ctx.session_id(), "session-789");
}
#[test]
fn test_context_with_branch() {
let agent = Arc::new(MockAgent { name: "test_agent".to_string() });
let content = Content::new("user");
let ctx = InvocationContext::new(
"inv-123".to_string(),
agent,
"user-456".to_string(),
"test-app".to_string(),
"session-789".to_string(),
content,
Arc::new(MockSessionWithState::new()),
)
.unwrap()
.with_branch("main.sub".to_string());
assert_eq!(ctx.branch(), "main.sub");
}
#[test]
fn test_context_with_run_config() {
let agent = Arc::new(MockAgent { name: "test_agent".to_string() });
let content = Content::new("user");
let config = RunConfig { streaming_mode: StreamingMode::SSE, ..RunConfig::default() };
let ctx = InvocationContext::new(
"inv-123".to_string(),
agent,
"user-456".to_string(),
"test-app".to_string(),
"session-789".to_string(),
content,
Arc::new(MockSessionWithState::new()),
)
.unwrap()
.with_run_config(config);
assert_eq!(ctx.run_config().streaming_mode, StreamingMode::SSE);
}
#[test]
fn test_context_end_invocation() {
let agent = Arc::new(MockAgent { name: "test_agent".to_string() });
let content = Content::new("user");
let ctx = InvocationContext::new(
"inv-123".to_string(),
agent,
"user-456".to_string(),
"test-app".to_string(),
"session-789".to_string(),
content,
Arc::new(MockSessionWithState::new()),
)
.unwrap();
assert!(!ctx.ended());
ctx.end_invocation();
assert!(ctx.ended());
}
#[test]
fn test_context_agent_access() {
let agent = Arc::new(MockAgent { name: "test_agent".to_string() });
let content = Content::new("user");
let ctx = InvocationContext::new(
"inv-123".to_string(),
agent.clone(),
"user-456".to_string(),
"test-app".to_string(),
"session-789".to_string(),
content,
Arc::new(MockSessionWithState::new()),
)
.unwrap();
let retrieved_agent = ctx.agent();
assert_eq!(retrieved_agent.name(), "test_agent");
}
#[test]
fn test_context_optional_services() {
let agent = Arc::new(MockAgent { name: "test_agent".to_string() });
let content = Content::new("user");
let ctx = InvocationContext::new(
"inv-123".to_string(),
agent,
"user-456".to_string(),
"test-app".to_string(),
"session-789".to_string(),
content,
Arc::new(MockSessionWithState::new()),
)
.unwrap();
assert!(ctx.artifacts().is_none());
assert!(ctx.memory().is_none());
}
#[test]
fn test_mutable_session_state_delta_propagation() {
let mut initial_state = HashMap::new();
initial_state.insert("initial_key".to_string(), serde_json::json!("initial_value"));
let session = Arc::new(MockSessionWithState::with_state(initial_state));
let mutable = MutableSession::new(session);
assert_eq!(mutable.state().get("initial_key"), Some(serde_json::json!("initial_value")));
let mut delta = HashMap::new();
delta.insert("research_findings".to_string(), serde_json::json!("AI research results"));
delta.insert("another_key".to_string(), serde_json::json!(42));
mutable.apply_state_delta(&delta);
assert_eq!(
mutable.state().get("research_findings"),
Some(serde_json::json!("AI research results"))
);
assert_eq!(mutable.state().get("another_key"), Some(serde_json::json!(42)));
assert_eq!(mutable.state().get("initial_key"), Some(serde_json::json!("initial_value")));
}
#[test]
fn test_mutable_session_temp_keys_not_persisted() {
let session = Arc::new(MockSessionWithState::new());
let mutable = MutableSession::new(session);
let mut delta = HashMap::new();
delta.insert("temp:scratch".to_string(), serde_json::json!("temporary"));
delta.insert("permanent".to_string(), serde_json::json!("persisted"));
mutable.apply_state_delta(&delta);
assert_eq!(mutable.state().get("temp:scratch"), None);
assert_eq!(mutable.state().get("permanent"), Some(serde_json::json!("persisted")));
}
#[test]
fn test_mutable_session_shared_across_contexts() {
let agent = Arc::new(MockAgent { name: "test_agent".to_string() });
let content = Content::new("user");
let ctx1 = InvocationContext::new(
"inv-1".to_string(),
agent.clone(),
"user-456".to_string(),
"test-app".to_string(),
"session-789".to_string(),
content.clone(),
Arc::new(MockSessionWithState::new()),
)
.unwrap();
let mut delta = HashMap::new();
delta.insert("from_agent1".to_string(), serde_json::json!("value_from_agent1"));
ctx1.mutable_session().apply_state_delta(&delta);
let ctx2 = InvocationContext::with_mutable_session(
"inv-2".to_string(),
agent.clone(),
"user-456".to_string(),
"test-app".to_string(),
"session-789".to_string(),
content,
ctx1.mutable_session().clone(),
)
.unwrap();
assert_eq!(
ctx2.session().state().get("from_agent1"),
Some(serde_json::json!("value_from_agent1"))
);
let mut delta2 = HashMap::new();
delta2.insert("from_agent2".to_string(), serde_json::json!("value_from_agent2"));
ctx2.mutable_session().apply_state_delta(&delta2);
assert_eq!(
ctx1.session().state().get("from_agent2"),
Some(serde_json::json!("value_from_agent2"))
);
}
#[test]
fn test_mutable_session_event_accumulation() {
let session = Arc::new(MockSessionWithState::new());
let mutable = MutableSession::new(session);
assert_eq!(mutable.conversation_history().len(), 0);
let mut event1 = Event::new("inv-1");
event1.author = "user".to_string();
event1.llm_response.content = Some(Content {
role: "user".to_string(),
parts: vec![Part::Text { text: "Hello".to_string() }],
});
mutable.append_event(event1);
let mut event2 = Event::new("inv-1");
event2.author = "assistant".to_string();
event2.llm_response.content = Some(Content {
role: "model".to_string(),
parts: vec![Part::Text { text: "Hi there!".to_string() }],
});
mutable.append_event(event2);
assert_eq!(mutable.events_len(), 2);
let history = mutable.conversation_history();
assert_eq!(history.len(), 2);
assert_eq!(history[0].role, "user");
assert_eq!(history[1].role, "model");
}
#[test]
fn test_mutable_session_state_all() {
let mut initial_state = HashMap::new();
initial_state.insert("key1".to_string(), serde_json::json!("value1"));
let session = Arc::new(MockSessionWithState::with_state(initial_state));
let mutable = MutableSession::new(session);
let mut delta = HashMap::new();
delta.insert("key2".to_string(), serde_json::json!("value2"));
mutable.apply_state_delta(&delta);
let all_state = mutable.state().all();
assert_eq!(all_state.len(), 2);
assert_eq!(all_state.get("key1"), Some(&serde_json::json!("value1")));
assert_eq!(all_state.get("key2"), Some(&serde_json::json!("value2")));
}
#[test]
fn conversation_history_preserves_tool_role() {
let session = Arc::new(MockSessionWithState::new());
let mutable = MutableSession::new(session);
let mut user_event = Event::new("inv-1");
user_event.author = "user".to_string();
user_event.llm_response.content = Some(Content {
role: "user".to_string(),
parts: vec![Part::Text { text: "hello".into() }],
});
mutable.append_event(user_event);
let mut assistant_event = Event::new("inv-1");
assistant_event.author = "my_agent".to_string();
assistant_event.llm_response.content = Some(Content {
role: "model".to_string(),
parts: vec![Part::FunctionCall {
name: "browser_navigate".into(),
args: serde_json::json!({"url": "https://example.com"}),
id: Some("call_1".into()),
thought_signature: None,
}],
});
mutable.append_event(assistant_event);
let mut tool_event = Event::new("inv-1");
tool_event.author = "my_agent".to_string();
tool_event.llm_response.content = Some(Content {
role: "function".to_string(),
parts: vec![Part::FunctionResponse {
function_response: FunctionResponseData {
name: "browser_navigate".into(),
response: serde_json::json!({"success": true}),
},
id: Some("call_1".into()),
}],
});
mutable.append_event(tool_event);
let history = mutable.conversation_history();
assert_eq!(history.len(), 3);
assert_eq!(history[0].role, "user");
assert_eq!(history[1].role, "model");
assert_eq!(history[2].role, "function"); }
#[test]
fn conversation_history_maps_agent_events_to_model() {
let session = Arc::new(MockSessionWithState::new());
let mutable = MutableSession::new(session);
let mut event = Event::new("inv-1");
event.author = "my_agent".to_string();
event.llm_response.content = Some(Content {
role: "model".to_string(),
parts: vec![Part::Text { text: "here are the results".into() }],
});
mutable.append_event(event);
let history = mutable.conversation_history();
assert_eq!(history[0].role, "model");
}
#[test]
fn conversation_history_for_agent_filters_other_agents_events() {
let session = Arc::new(MockSessionWithState::new());
let mutable = MutableSession::new(session);
let mut user_event = Event::new("inv-1");
user_event.author = "user".to_string();
user_event.llm_response.content = Some(Content {
role: "user".to_string(),
parts: vec![Part::Text { text: "find me a job".into() }],
});
mutable.append_event(user_event);
let mut coord_call = Event::new("inv-1");
coord_call.author = "coordinator".to_string();
coord_call.llm_response.content = Some(Content {
role: "model".to_string(),
parts: vec![Part::FunctionCall {
name: "get_profile".into(),
args: serde_json::json!({}),
id: Some("call_1".into()),
thought_signature: None,
}],
});
mutable.append_event(coord_call);
let mut coord_tool_resp = Event::new("inv-1");
coord_tool_resp.author = "coordinator".to_string();
coord_tool_resp.llm_response.content = Some(Content {
role: "function".to_string(),
parts: vec![Part::FunctionResponse {
function_response: FunctionResponseData {
name: "get_profile".into(),
response: serde_json::json!({"name": "Alice"}),
},
id: Some("call_1".into()),
}],
});
mutable.append_event(coord_tool_resp);
let mut transfer_event = Event::new("inv-1");
transfer_event.author = "coordinator".to_string();
transfer_event.actions.transfer_to_agent = Some("sourcing_agent".to_string());
transfer_event.llm_response.content = Some(Content {
role: "model".to_string(),
parts: vec![Part::Text { text: "transferring".into() }],
});
mutable.append_event(transfer_event);
let full_history = mutable.conversation_history();
assert_eq!(full_history.len(), 4);
use adk_core::Session;
let filtered = mutable.conversation_history_for_agent("sourcing_agent");
assert_eq!(filtered.len(), 1);
assert_eq!(filtered[0].role, "user");
}
#[test]
fn conversation_history_for_agent_keeps_own_events() {
let session = Arc::new(MockSessionWithState::new());
let mutable = MutableSession::new(session);
let mut user_event = Event::new("inv-1");
user_event.author = "user".to_string();
user_event.llm_response.content = Some(Content {
role: "user".to_string(),
parts: vec![Part::Text { text: "search for jobs".into() }],
});
mutable.append_event(user_event);
let mut agent_call = Event::new("inv-2");
agent_call.author = "sourcing_agent".to_string();
agent_call.llm_response.content = Some(Content {
role: "model".to_string(),
parts: vec![Part::FunctionCall {
name: "search_jobs".into(),
args: serde_json::json!({"query": "rust developer"}),
id: Some("call_2".into()),
thought_signature: None,
}],
});
mutable.append_event(agent_call);
let mut tool_resp = Event::new("inv-2");
tool_resp.author = "sourcing_agent".to_string();
tool_resp.llm_response.content = Some(Content {
role: "function".to_string(),
parts: vec![Part::FunctionResponse {
function_response: FunctionResponseData {
name: "search_jobs".into(),
response: serde_json::json!({"jobs": []}),
},
id: Some("call_2".into()),
}],
});
mutable.append_event(tool_resp);
use adk_core::Session;
let filtered = mutable.conversation_history_for_agent("sourcing_agent");
assert_eq!(filtered.len(), 3);
assert_eq!(filtered[0].role, "user");
assert_eq!(filtered[1].role, "model");
assert_eq!(filtered[2].role, "function");
}
#[test]
fn conversation_history_for_agent_excludes_other_agents_function_responses() {
let session = Arc::new(MockSessionWithState::new());
let mutable = MutableSession::new(session);
let mut user_event = Event::new("inv-1");
user_event.author = "user".to_string();
user_event.llm_response.content = Some(Content {
role: "user".to_string(),
parts: vec![Part::Text { text: "hello".into() }],
});
mutable.append_event(user_event);
let mut other_tool = Event::new("inv-1");
other_tool.author = "other_agent".to_string();
other_tool.llm_response.content = Some(Content {
role: "function".to_string(),
parts: vec![Part::FunctionResponse {
function_response: FunctionResponseData {
name: "some_tool".into(),
response: serde_json::json!({"data": "value"}),
},
id: None,
}],
});
mutable.append_event(other_tool);
let mut other_model = Event::new("inv-1");
other_model.author = "other_agent".to_string();
other_model.llm_response.content = Some(Content {
role: "model".to_string(),
parts: vec![Part::Text { text: "I did something".into() }],
});
mutable.append_event(other_model);
use adk_core::Session;
let filtered = mutable.conversation_history_for_agent("my_agent");
assert_eq!(filtered.len(), 1);
assert_eq!(filtered[0].role, "user");
}
#[test]
fn conversation_history_for_agent_double_transfer_sees_own_prior_events() {
let session = Arc::new(MockSessionWithState::new());
let mutable = MutableSession::new(session);
let mut user_event = Event::new("inv-1");
user_event.author = "user".to_string();
user_event.llm_response.content = Some(Content {
role: "user".to_string(),
parts: vec![Part::Text { text: "find me a rust job".into() }],
});
mutable.append_event(user_event);
let mut coord_analyze = Event::new("inv-1");
coord_analyze.author = "coordinator".to_string();
coord_analyze.llm_response.content = Some(Content {
role: "model".to_string(),
parts: vec![Part::Text { text: "I'll delegate to sourcing".into() }],
});
mutable.append_event(coord_analyze);
let mut transfer1 = Event::new("inv-1");
transfer1.author = "coordinator".to_string();
transfer1.actions.transfer_to_agent = Some("sourcing_agent".to_string());
mutable.append_event(transfer1);
let mut sourcing_call1 = Event::new("inv-2");
sourcing_call1.author = "sourcing_agent".to_string();
sourcing_call1.llm_response.content = Some(Content {
role: "model".to_string(),
parts: vec![Part::FunctionCall {
name: "search_jobs".into(),
args: serde_json::json!({"query": "rust developer"}),
id: Some("call_s1".into()),
thought_signature: None,
}],
});
mutable.append_event(sourcing_call1);
let mut sourcing_resp1 = Event::new("inv-2");
sourcing_resp1.author = "sourcing_agent".to_string();
sourcing_resp1.llm_response.content = Some(Content {
role: "function".to_string(),
parts: vec![Part::FunctionResponse {
function_response: FunctionResponseData {
name: "search_jobs".into(),
response: serde_json::json!({"jobs": ["job_1", "job_2"]}),
},
id: Some("call_s1".into()),
}],
});
mutable.append_event(sourcing_resp1);
let mut transfer_back = Event::new("inv-2");
transfer_back.author = "sourcing_agent".to_string();
transfer_back.actions.transfer_to_agent = Some("coordinator".to_string());
transfer_back.llm_response.content = Some(Content {
role: "model".to_string(),
parts: vec![Part::Text { text: "Found 2 jobs, transferring back".into() }],
});
mutable.append_event(transfer_back);
let mut coord_work = Event::new("inv-3");
coord_work.author = "coordinator".to_string();
coord_work.llm_response.content = Some(Content {
role: "model".to_string(),
parts: vec![Part::FunctionCall {
name: "rank_candidates".into(),
args: serde_json::json!({}),
id: Some("call_c2".into()),
thought_signature: None,
}],
});
mutable.append_event(coord_work);
let mut coord_resp = Event::new("inv-3");
coord_resp.author = "coordinator".to_string();
coord_resp.llm_response.content = Some(Content {
role: "function".to_string(),
parts: vec![Part::FunctionResponse {
function_response: FunctionResponseData {
name: "rank_candidates".into(),
response: serde_json::json!({"ranked": ["job_1"]}),
},
id: Some("call_c2".into()),
}],
});
mutable.append_event(coord_resp);
let mut transfer2 = Event::new("inv-3");
transfer2.author = "coordinator".to_string();
transfer2.actions.transfer_to_agent = Some("sourcing_agent".to_string());
mutable.append_event(transfer2);
use adk_core::Session;
let filtered = mutable.conversation_history_for_agent("sourcing_agent");
assert_eq!(
filtered.len(),
4,
"Expected 4 events (user + 3 own), got {}: {:?}",
filtered.len(),
filtered.iter().map(|c| (&c.role, c.parts.first())).collect::<Vec<_>>()
);
assert_eq!(filtered[0].role, "user");
assert_eq!(filtered[1].role, "model"); assert_eq!(filtered[2].role, "function"); assert_eq!(filtered[3].role, "model"); }