rskit_agent/agent/
mod.rs

1//! Agent — the multi-turn agentic execution loop.
2
3use std::sync::Arc;
4
5use rskit_llm::provider::Provider;
6
7use crate::config::AgentConfig;
8
9mod component;
10mod run;
11mod stream;
12
13/// A multi-turn agentic loop that drives an LLM provider, executes tool calls,
14/// and emits hook events at each lifecycle point.
15pub struct Agent {
16    provider: Arc<dyn Provider>,
17    config: AgentConfig,
18}
19
20impl Agent {
21    /// Create a new agent with the given provider and configuration.
22    pub fn new(provider: Arc<dyn Provider>, config: AgentConfig) -> Self {
23        Self { provider, config }
24    }
25
26    /// Create a new agent with the locked default configuration.
27    pub fn with_defaults(provider: Arc<dyn Provider>) -> Self {
28        Self::new(provider, AgentConfig::default())
29    }
30}
31
32#[cfg(test)]
33mod tests {
34    use super::*;
35    use async_trait::async_trait;
36    use std::pin::Pin;
37
38    use futures::{Stream, StreamExt};
39    use rskit_ai::Capabilities;
40    use rskit_ai::StreamEventRef;
41    use rskit_ai::chat::count_tokens_approx;
42    use rskit_errors::AppError;
43    use rskit_hook::{HookError, HookRegistry};
44    use rskit_llm::types::{
45        self, AssistantMessage, CompletionRequest, CompletionResponse, Message, Usage,
46    };
47    use rskit_resilience::{ConstantBackoff, Policy, RetryPolicy};
48    use rskit_tool::{Context, Registry};
49    use std::sync::atomic::{AtomicU32, Ordering};
50    use std::time::Duration;
51
52    use crate::types::{AgentEvent, StopReason};
53
54    // ── Mock provider ───────────────────────────────────────────────────
55
56    struct MockProvider {
57        responses: Vec<CompletionResponse>,
58        call_count: AtomicU32,
59    }
60
61    impl MockProvider {
62        fn new(responses: Vec<CompletionResponse>) -> Self {
63            Self {
64                responses,
65                call_count: AtomicU32::new(0),
66            }
67        }
68
69        fn single_text(text: &str) -> Self {
70            Self::new(vec![CompletionResponse {
71                message: AssistantMessage {
72                    content: types::text_content(text),
73                    tool_calls: vec![],
74                    usage: None,
75                },
76                model: "mock".to_string(),
77                usage: Usage {
78                    input_tokens: 10,
79                    output_tokens: 5,
80                    cached_tokens: 0,
81                    reasoning_tokens: 0,
82                },
83                stop_reason: Some(rskit_llm::FinishReason::Stop),
84            }])
85        }
86    }
87
88    #[async_trait]
89    impl rskit_provider::Provider for MockProvider {
90        fn name(&self) -> &'static str {
91            "mock"
92        }
93    }
94
95    #[async_trait]
96    impl rskit_provider::RequestResponse<CompletionRequest, CompletionResponse> for MockProvider {
97        async fn execute(&self, input: CompletionRequest) -> Result<CompletionResponse, AppError> {
98            self.complete(input).await
99        }
100    }
101
102    #[async_trait]
103    impl Provider for MockProvider {
104        async fn complete(
105            &self,
106            _request: CompletionRequest,
107        ) -> Result<CompletionResponse, AppError> {
108            let idx = self.call_count.fetch_add(1, Ordering::SeqCst) as usize;
109            if idx < self.responses.len() {
110                Ok(self.responses[idx].clone())
111            } else {
112                // Return last response for any additional calls
113                Ok(self.responses.last().unwrap().clone())
114            }
115        }
116
117        async fn stream(
118            &self,
119            _request: CompletionRequest,
120        ) -> Result<Pin<Box<dyn Stream<Item = StreamEventRef> + Send>>, AppError> {
121            Ok(Box::pin(futures::stream::empty()))
122        }
123
124        fn capabilities(&self) -> Capabilities {
125            Capabilities {
126                tool_use: true,
127                streaming: false,
128                max_input_tokens: Some(128_000),
129                max_output_tokens: Some(4_096),
130                ..Default::default()
131            }
132        }
133
134        fn count_tokens(&self, messages: &[Message]) -> usize {
135            count_tokens_approx(messages)
136        }
137    }
138
139    // ── Tests ───────────────────────────────────────────────────────────
140
141    #[tokio::test]
142    async fn test_agent_simple_completion() {
143        let provider = Arc::new(MockProvider::single_text("Hello!"));
144        let agent = Agent::new(
145            provider,
146            AgentConfig {
147                tools: None,
148                hooks: None,
149                system_prompt: "You are helpful.".to_string(),
150                max_turns: 5,
151                max_tokens: 100_000,
152                wall_clock: Duration::from_mins(1),
153                max_tool_calls: 50,
154                tool_concurrency: 4,
155                tool_timeout: Duration::from_secs(30),
156                policy: None,
157                context_strategy: None,
158                model: String::new(),
159            },
160        );
161
162        let result = agent.run(vec![types::user("Hi")]).await.unwrap();
163        assert_eq!(result.turn_count, 1);
164        assert!(matches!(result.stop_reason, StopReason::EndTurn));
165        assert_eq!(result.total_usage.input_tokens, 10);
166        assert_eq!(result.total_usage.output_tokens, 5);
167    }
168
169    #[tokio::test]
170    async fn test_agent_max_turns() {
171        // Provider always returns tool calls → agent loops until max_turns
172        let tool_call_response = CompletionResponse {
173            message: AssistantMessage {
174                content: vec![],
175                tool_calls: vec![rskit_llm::ToolUseBlock {
176                    id: "tc_1".to_string(),
177                    name: "test_tool".to_string(),
178                    input: serde_json::json!({"x": 1}).as_object().cloned().unwrap(),
179                }],
180                usage: None,
181            },
182            model: "mock".to_string(),
183            usage: Usage {
184                input_tokens: 5,
185                output_tokens: 5,
186                cached_tokens: 0,
187                reasoning_tokens: 0,
188            },
189            stop_reason: Some(rskit_llm::FinishReason::ToolUse),
190        };
191
192        let provider = Arc::new(MockProvider::new(vec![tool_call_response]));
193
194        // No tools registered → tool calls will fail, but loop continues
195        let agent = Agent::new(
196            provider,
197            AgentConfig {
198                tools: None,
199                hooks: None,
200                system_prompt: "sys".to_string(),
201                max_turns: 3,
202                max_tokens: 100_000,
203                wall_clock: Duration::from_mins(1),
204                max_tool_calls: 50,
205                tool_concurrency: 4,
206                tool_timeout: Duration::from_secs(30),
207                policy: None,
208                context_strategy: None,
209                model: String::new(),
210            },
211        );
212
213        let result = agent.run(vec![types::user("go")]).await.unwrap();
214        assert_eq!(result.turn_count, 3);
215        assert!(matches!(result.stop_reason, StopReason::MaxTurns));
216    }
217
218    #[tokio::test]
219    async fn test_agent_with_tool() {
220        use rskit_tool::{from_fn, text_result};
221        use schemars::JsonSchema;
222        use serde::Deserialize;
223
224        #[derive(Deserialize, JsonSchema)]
225        struct AddInput {
226            a: i32,
227            b: i32,
228        }
229
230        let registry = Arc::new(Registry::new());
231        let attempts = Arc::new(AtomicU32::new(0));
232        let attempts_for_tool = Arc::clone(&attempts);
233        registry
234            .register(
235                from_fn(
236                    "add",
237                    "Add two numbers",
238                    move |_ctx: Context, input: AddInput| {
239                        let attempts = Arc::clone(&attempts_for_tool);
240                        async move {
241                            let attempt = attempts.fetch_add(1, Ordering::SeqCst);
242                            if attempt == 0 {
243                                Err(AppError::connection_failed("add"))
244                            } else {
245                                Ok(text_result(&format!("{}", input.a + input.b)))
246                            }
247                        }
248                    },
249                )
250                .unwrap(),
251            )
252            .unwrap();
253
254        // First call: model requests tool
255        let tool_call_resp = CompletionResponse {
256            message: AssistantMessage {
257                content: vec![],
258                tool_calls: vec![rskit_llm::ToolUseBlock {
259                    id: "tc_1".to_string(),
260                    name: "add".to_string(),
261                    input: serde_json::json!({"a": 2, "b": 3})
262                        .as_object()
263                        .cloned()
264                        .unwrap(),
265                }],
266                usage: None,
267            },
268            model: "mock".to_string(),
269            usage: Usage {
270                input_tokens: 10,
271                output_tokens: 5,
272                cached_tokens: 0,
273                reasoning_tokens: 0,
274            },
275            stop_reason: Some(rskit_llm::FinishReason::ToolUse),
276        };
277
278        // Second call: model returns final text
279        let final_resp = CompletionResponse {
280            message: AssistantMessage {
281                content: types::text_content("The answer is 5"),
282                tool_calls: vec![],
283                usage: None,
284            },
285            model: "mock".to_string(),
286            usage: Usage {
287                input_tokens: 15,
288                output_tokens: 8,
289                cached_tokens: 0,
290                reasoning_tokens: 0,
291            },
292            stop_reason: Some(rskit_llm::FinishReason::Stop),
293        };
294
295        let provider = Arc::new(MockProvider::new(vec![tool_call_resp, final_resp]));
296
297        let agent = Agent::new(
298            provider,
299            AgentConfig {
300                tools: Some(registry),
301                hooks: None,
302                system_prompt: "You are a calculator.".to_string(),
303                max_turns: 5,
304                max_tokens: 100_000,
305                wall_clock: Duration::from_mins(1),
306                max_tool_calls: 50,
307                tool_concurrency: 4,
308                tool_timeout: Duration::from_secs(30),
309                policy: Some(
310                    Policy::new().with_retry(
311                        RetryPolicy::new()
312                            .with_max_attempts(2)
313                            .with_constant_backoff(ConstantBackoff::new(Duration::from_millis(1)))
314                            .with_jitter(false),
315                    ),
316                ),
317                context_strategy: None,
318                model: String::new(),
319            },
320        );
321
322        let result = agent.run(vec![types::user("What is 2+3?")]).await.unwrap();
323        assert_eq!(attempts.load(Ordering::SeqCst), 2);
324        assert_eq!(result.turn_count, 2);
325        assert!(matches!(result.stop_reason, StopReason::EndTurn));
326        // Usage should be summed
327        assert_eq!(result.total_usage.input_tokens, 25);
328        assert_eq!(result.total_usage.output_tokens, 13);
329    }
330
331    #[tokio::test]
332    async fn test_agent_max_budget() {
333        let tool_call_response = CompletionResponse {
334            message: AssistantMessage {
335                content: vec![],
336                tool_calls: vec![rskit_llm::ToolUseBlock {
337                    id: "tc_1".to_string(),
338                    name: "noop".to_string(),
339                    input: serde_json::Map::new(),
340                }],
341                usage: None,
342            },
343            model: "mock".to_string(),
344            usage: Usage {
345                input_tokens: 50,
346                output_tokens: 50,
347                cached_tokens: 0,
348                reasoning_tokens: 0,
349            },
350            stop_reason: Some(rskit_llm::FinishReason::ToolUse),
351        };
352
353        let provider = Arc::new(MockProvider::new(vec![tool_call_response]));
354
355        let agent = Agent::new(
356            provider,
357            AgentConfig {
358                tools: None,
359                hooks: None,
360                system_prompt: "sys".to_string(),
361                max_turns: 100,
362                max_tokens: 80, // Budget of 80 tokens total
363                wall_clock: Duration::from_mins(1),
364                max_tool_calls: 50,
365                tool_concurrency: 4,
366                tool_timeout: Duration::from_secs(30),
367                policy: None,
368                context_strategy: None,
369                model: String::new(),
370            },
371        );
372
373        let result = agent.run(vec![types::user("go")]).await.unwrap();
374        assert!(matches!(result.stop_reason, StopReason::MaxTokens));
375    }
376
377    #[tokio::test]
378    async fn max_budget_stops_after_final_response_without_tool_calls() {
379        let provider = Arc::new(MockProvider::new(vec![CompletionResponse {
380            message: AssistantMessage {
381                content: types::text_content("large final response"),
382                tool_calls: vec![],
383                usage: None,
384            },
385            model: "mock".to_string(),
386            usage: Usage {
387                input_tokens: 50,
388                output_tokens: 50,
389                cached_tokens: 0,
390                reasoning_tokens: 0,
391            },
392            stop_reason: Some(rskit_llm::FinishReason::Stop),
393        }]));
394
395        let agent = Agent::new(
396            provider,
397            AgentConfig {
398                tools: None,
399                hooks: None,
400                system_prompt: "sys".to_string(),
401                max_turns: 5,
402                max_tokens: 80,
403                wall_clock: Duration::from_mins(1),
404                max_tool_calls: 50,
405                tool_concurrency: 4,
406                tool_timeout: Duration::from_secs(30),
407                policy: None,
408                context_strategy: None,
409                model: String::new(),
410            },
411        );
412
413        let result = agent.run(vec![types::user("go")]).await.unwrap();
414        assert!(matches!(result.stop_reason, StopReason::MaxTokens));
415        assert_eq!(result.turn_count, 1);
416    }
417
418    #[tokio::test]
419    async fn test_agent_hook_fatal_error_stops() {
420        let provider = Arc::new(MockProvider::single_text("Hello"));
421        let hooks = Arc::new(HookRegistry::new());
422
423        let _unsub = hooks.on::<crate::hooks::TurnStart>(crate::turn_start_type(), |_, _| {
424            Err(HookError::fatal("blocked by policy"))
425        });
426
427        let agent = Agent::new(
428            provider,
429            AgentConfig {
430                tools: None,
431                hooks: Some(hooks),
432                system_prompt: "sys".to_string(),
433                max_turns: 5,
434                max_tokens: 100_000,
435                wall_clock: Duration::from_mins(1),
436                max_tool_calls: 50,
437                tool_concurrency: 4,
438                tool_timeout: Duration::from_secs(30),
439                policy: None,
440                context_strategy: None,
441                model: String::new(),
442            },
443        );
444
445        let result = agent.run(vec![types::user("hi")]).await.unwrap();
446        assert!(matches!(result.stop_reason, StopReason::Aborted));
447        assert_eq!(result.turn_count, 0);
448    }
449
450    #[tokio::test]
451    async fn hook_observes_request_without_mutation_surface() {
452        let provider = Arc::new(MockProvider::single_text("done"));
453        let hooks = Arc::new(HookRegistry::new());
454        let observed = Arc::new(AtomicU32::new(0));
455        let observed_clone = Arc::clone(&observed);
456
457        let _unsub =
458            hooks.on::<crate::hooks::PreLLMCall>(crate::pre_llm_call_type(), move |_, event| {
459                assert_eq!(event.request.model, "test-model");
460                observed_clone.fetch_add(1, Ordering::SeqCst);
461                Ok(())
462            });
463
464        let agent = Agent::new(
465            provider,
466            AgentConfig {
467                hooks: Some(hooks),
468                system_prompt: "sys".to_string(),
469                model: "test-model".to_string(),
470                ..AgentConfig::default()
471            },
472        );
473
474        let result = agent.run(vec![types::user("hi")]).await.unwrap();
475        assert!(matches!(result.stop_reason, StopReason::EndTurn));
476        assert_eq!(observed.load(Ordering::SeqCst), 1);
477    }
478
479    #[tokio::test]
480    async fn test_agent_hook_counts() {
481        let provider = Arc::new(MockProvider::single_text("done"));
482        let hooks = Arc::new(HookRegistry::new());
483
484        let pre_count = Arc::new(AtomicU32::new(0));
485        let post_count = Arc::new(AtomicU32::new(0));
486
487        let pc = pre_count.clone();
488        let _unsub1 =
489            hooks.on::<crate::hooks::PreLLMCall>(crate::pre_llm_call_type(), move |_, _| {
490                pc.fetch_add(1, Ordering::SeqCst);
491                Ok(())
492            });
493
494        let poc = post_count.clone();
495        let _unsub2 =
496            hooks.on::<crate::hooks::PostLLMCall>(crate::post_llm_call_type(), move |_, _| {
497                poc.fetch_add(1, Ordering::SeqCst);
498                Ok(())
499            });
500
501        let agent = Agent::new(
502            provider,
503            AgentConfig {
504                tools: None,
505                hooks: Some(hooks),
506                system_prompt: "sys".to_string(),
507                max_turns: 5,
508                max_tokens: 100_000,
509                wall_clock: Duration::from_mins(1),
510                max_tool_calls: 50,
511                tool_concurrency: 4,
512                tool_timeout: Duration::from_secs(30),
513                policy: None,
514                context_strategy: None,
515                model: String::new(),
516            },
517        );
518
519        agent.run(vec![types::user("hi")]).await.unwrap();
520        assert_eq!(pre_count.load(Ordering::SeqCst), 1);
521        assert_eq!(post_count.load(Ordering::SeqCst), 1);
522    }
523
524    #[tokio::test]
525    async fn test_agent_stream() {
526        let provider = Arc::new(MockProvider::single_text("streamed"));
527        let agent = Agent::new(
528            provider,
529            AgentConfig {
530                tools: None,
531                hooks: None,
532                system_prompt: "sys".to_string(),
533                max_turns: 5,
534                max_tokens: 100_000,
535                wall_clock: Duration::from_mins(1),
536                max_tool_calls: 50,
537                tool_concurrency: 4,
538                tool_timeout: Duration::from_secs(30),
539                policy: None,
540                context_strategy: None,
541                model: String::new(),
542            },
543        );
544
545        let stream = agent.stream(vec![types::user("hi")]);
546        let events: Vec<AgentEvent> = stream.collect().await;
547        assert_eq!(events.len(), 3);
548        assert!(matches!(
549            events.first(),
550            Some(AgentEvent::TurnStart { turn: 0 })
551        ));
552        assert!(matches!(
553            events.get(1),
554            Some(AgentEvent::TurnComplete {
555                turn: 0,
556                usage: Usage {
557                    input_tokens: 10,
558                    output_tokens: 5,
559                    cached_tokens: 0,
560                    reasoning_tokens: 0,
561                },
562                ..
563            })
564        ));
565        assert!(matches!(events.last(), Some(AgentEvent::Complete { .. })));
566    }
567
568    #[tokio::test]
569    async fn agent_component_lifecycle_is_named_and_healthy() {
570        use rskit_component::Component;
571
572        let agent = Agent::with_defaults(Arc::new(MockProvider::single_text("ok")));
573
574        assert_eq!(Component::name(&agent), "rskit-agent");
575        agent.start().await.expect("start is a no-op success");
576        agent.stop().await.expect("stop is a no-op success");
577
578        let health = agent.health();
579        assert!(health.is_healthy());
580    }
581}
rskit_agent/agent/mod.rs

rskit_agent/agent/
mod.rs
