tirea_agentos_server/protocol/ag_ui/
runtime.rs

1//! Runtime wiring for AG-UI requests.
2//!
3//! Applies AG-UI–specific extensions to a [`ResolvedRun`]:
4//! frontend tool descriptor stubs, frontend suspended-call strategy,
5//! and context injection.
6
7use async_trait::async_trait;
8use serde_json::Value;
9use std::collections::HashSet;
10use std::sync::Arc;
11use tirea_agentos::runtime::loop_runner::{
12    BaseAgent, ParallelToolExecutor, SequentialToolExecutor,
13};
14use tirea_agentos::runtime::{compose_behaviors, ResolvedRun};
15use tirea_contract::runtime::behavior::ReadOnlyContext;
16use tirea_contract::runtime::phase::{
17    ActionSet, BeforeInferenceAction, BeforeToolExecuteAction, SuspendTicket,
18};
19use tirea_contract::runtime::tool_call::{Tool, ToolDescriptor, ToolError, ToolResult};
20use tirea_contract::runtime::AgentBehavior;
21use tirea_contract::runtime::{PendingToolCall, ToolCallResumeMode};
22use tirea_contract::ToolCallContext;
23
24use tirea_protocol_ag_ui::{build_context_addendum, RunAgentInput};
25
26/// Apply AG-UI–specific extensions to a [`ResolvedRun`].
27///
28/// Injects frontend tool stubs, suspended-call plugins, context
29/// injection, and request model/config overrides.
30pub fn apply_agui_extensions(resolved: &mut ResolvedRun, request: &RunAgentInput) {
31    if let Some(model) = request.model.as_ref().filter(|m| !m.trim().is_empty()) {
32        resolved.agent.model = model.clone();
33    }
34    if let Some(system_prompt) = request
35        .system_prompt
36        .as_ref()
37        .filter(|prompt| !prompt.trim().is_empty())
38    {
39        resolved.agent.system_prompt = system_prompt.clone();
40    }
41    if let Some(config) = request.config.clone() {
42        apply_agui_tool_execution_mode_override(resolved, &config);
43        apply_agui_chat_options_overrides(resolved, &config);
44    }
45
46    let frontend_defs = request.frontend_tools();
47    let frontend_tool_names: HashSet<String> =
48        frontend_defs.iter().map(|tool| tool.name.clone()).collect();
49
50    // Frontend tools → insert into resolved.tools (overlay semantics)
51    for tool in frontend_defs {
52        let stub = Arc::new(FrontendToolStub::new(
53            tool.name.clone(),
54            tool.description.clone(),
55            tool.parameters.clone(),
56        ));
57        let id = stub.descriptor().id.clone();
58        resolved.tools.entry(id).or_insert(stub as Arc<dyn Tool>);
59    }
60
61    // Run-scoped behaviors
62    if !frontend_tool_names.is_empty() {
63        add_behavior_mut(
64            &mut resolved.agent,
65            Arc::new(FrontendToolPendingPlugin::new(frontend_tool_names)),
66        );
67    }
68
69    // Context injection: forward useCopilotReadable context to the agent's system prompt.
70    if let Some(addendum) = build_context_addendum(request) {
71        add_behavior_mut(
72            &mut resolved.agent,
73            Arc::new(ContextInjectionPlugin::new(addendum)),
74        );
75    }
76}
77
78/// Add a behavior to a `BaseAgent` by reference, composing with any existing behavior.
79fn add_behavior_mut(agent: &mut BaseAgent, behavior: Arc<dyn AgentBehavior>) {
80    if agent.behavior.id() == "noop" {
81        agent.behavior = behavior;
82    } else {
83        let id = format!("{}+{}", agent.behavior.id(), behavior.id());
84        agent.behavior = compose_behaviors(id, vec![agent.behavior.clone(), behavior]);
85    }
86}
87
88fn apply_agui_tool_execution_mode_override(resolved: &mut ResolvedRun, config: &Value) {
89    let mode = config
90        .get("toolExecutionMode")
91        .or_else(|| config.get("tool_execution_mode"))
92        .and_then(Value::as_str)
93        .map(|value| value.trim().to_ascii_lowercase());
94
95    match mode.as_deref() {
96        Some("sequential") => {
97            resolved.agent.tool_executor = Arc::new(SequentialToolExecutor);
98        }
99        Some("parallel_batch_approval") => {
100            resolved.agent.tool_executor = Arc::new(ParallelToolExecutor::batch_approval());
101        }
102        Some("parallel_streaming") => {
103            resolved.agent.tool_executor = Arc::new(ParallelToolExecutor::streaming());
104        }
105        _ => {}
106    }
107}
108
109fn apply_agui_chat_options_overrides(resolved: &mut ResolvedRun, config: &Value) {
110    let mut chat_options = resolved.agent.chat_options.clone().unwrap_or_default();
111    let mut changed = false;
112
113    if let Some(map) = config.as_object() {
114        if let Some(value) = get_bool(map, "captureReasoningContent", "capture_reasoning_content") {
115            chat_options.capture_reasoning_content = Some(value);
116            changed = true;
117        }
118        if let Some(value) = get_bool(
119            map,
120            "normalizeReasoningContent",
121            "normalize_reasoning_content",
122        ) {
123            chat_options.normalize_reasoning_content = Some(value);
124            changed = true;
125        }
126    }
127
128    if let Some(map) = config
129        .get("chatOptions")
130        .and_then(Value::as_object)
131        .or_else(|| config.get("chat_options").and_then(Value::as_object))
132    {
133        if let Some(value) = get_bool(map, "captureReasoningContent", "capture_reasoning_content") {
134            chat_options.capture_reasoning_content = Some(value);
135            changed = true;
136        }
137        if let Some(value) = get_bool(
138            map,
139            "normalizeReasoningContent",
140            "normalize_reasoning_content",
141        ) {
142            chat_options.normalize_reasoning_content = Some(value);
143            changed = true;
144        }
145    }
146
147    if changed {
148        resolved.agent.chat_options = Some(chat_options);
149    }
150}
151
152fn get_bool(map: &serde_json::Map<String, Value>, primary: &str, alias: &str) -> Option<bool> {
153    map.get(primary)
154        .or_else(|| map.get(alias))
155        .and_then(Value::as_bool)
156}
157
158/// Runtime-only frontend tool descriptor stub.
159///
160/// The frontend pending plugin intercepts configured frontend tools before
161/// backend execution. This stub exists only to expose tool descriptors to the model.
162struct FrontendToolStub {
163    descriptor: ToolDescriptor,
164}
165
166impl FrontendToolStub {
167    fn new(name: String, description: String, parameters: Option<Value>) -> Self {
168        let mut descriptor = ToolDescriptor::new(&name, &name, description);
169        if let Some(parameters) = parameters {
170            descriptor = descriptor.with_parameters(parameters);
171        }
172        Self { descriptor }
173    }
174}
175
176#[async_trait]
177impl Tool for FrontendToolStub {
178    fn descriptor(&self) -> ToolDescriptor {
179        self.descriptor.clone()
180    }
181
182    async fn execute(
183        &self,
184        _args: Value,
185        _ctx: &ToolCallContext<'_>,
186    ) -> Result<ToolResult, ToolError> {
187        Ok(ToolResult::error(
188            &self.descriptor.id,
189            "frontend tool stub should be intercepted before backend execution",
190        ))
191    }
192}
193
194/// Run-scoped plugin that injects AG-UI context (from `useCopilotReadable`)
195/// into the agent's system prompt before inference.
196struct ContextInjectionPlugin {
197    addendum: String,
198}
199
200impl ContextInjectionPlugin {
201    fn new(addendum: String) -> Self {
202        Self { addendum }
203    }
204}
205
206#[async_trait]
207impl AgentBehavior for ContextInjectionPlugin {
208    fn id(&self) -> &str {
209        "agui_context_injection"
210    }
211
212    async fn before_inference(
213        &self,
214        _ctx: &ReadOnlyContext<'_>,
215    ) -> ActionSet<BeforeInferenceAction> {
216        ActionSet::single(BeforeInferenceAction::AddSystemContext(
217            self.addendum.clone(),
218        ))
219    }
220}
221
222/// Run-scoped frontend interaction strategy for AG-UI.
223struct FrontendToolPendingPlugin {
224    frontend_tools: HashSet<String>,
225}
226
227impl FrontendToolPendingPlugin {
228    fn new(frontend_tools: HashSet<String>) -> Self {
229        Self { frontend_tools }
230    }
231}
232
233#[async_trait]
234impl AgentBehavior for FrontendToolPendingPlugin {
235    fn id(&self) -> &str {
236        "agui_frontend_tools"
237    }
238
239    async fn before_tool_execute(
240        &self,
241        ctx: &ReadOnlyContext<'_>,
242    ) -> ActionSet<BeforeToolExecuteAction> {
243        let Some(tool_name) = ctx.tool_name() else {
244            return ActionSet::empty();
245        };
246        if !self.frontend_tools.contains(tool_name) {
247            return ActionSet::empty();
248        }
249
250        if let Some(resume) = ctx.resume_input() {
251            let result = match resume.action {
252                tirea_contract::io::ResumeDecisionAction::Resume => {
253                    ToolResult::success(tool_name.to_string(), resume.result.clone())
254                }
255                tirea_contract::io::ResumeDecisionAction::Cancel => ToolResult::error(
256                    tool_name.to_string(),
257                    resume
258                        .reason
259                        .clone()
260                        .filter(|r| !r.trim().is_empty())
261                        .unwrap_or_else(|| "User denied the action".to_string()),
262                ),
263            };
264            return ActionSet::single(BeforeToolExecuteAction::SetToolResult(result));
265        }
266        let Some(call_id) = ctx.tool_call_id().map(str::to_string) else {
267            return ActionSet::empty();
268        };
269
270        let args = ctx.tool_args().cloned().unwrap_or_default();
271        let suspension = tirea_contract::Suspension::new(&call_id, format!("tool:{tool_name}"))
272            .with_parameters(args.clone());
273        ActionSet::single(BeforeToolExecuteAction::Suspend(SuspendTicket::new(
274            suspension,
275            PendingToolCall::new(call_id, tool_name.to_string(), args),
276            ToolCallResumeMode::UseDecisionAsToolResult,
277        )))
278    }
279}
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284    use async_trait::async_trait;
285    use serde_json::json;
286    use std::collections::HashMap;
287    use std::sync::Arc;
288    use tirea_agentos::runtime::loop_runner::BaseAgent;
289    use tirea_contract::runtime::phase::{Phase, StepContext};
290    use tirea_contract::runtime::tool_call::ToolGate;
291    use tirea_contract::testing::TestFixture;
292    use tirea_contract::thread::ToolCall;
293    use tirea_protocol_ag_ui::{Context, Message, ToolExecutionLocation};
294
295    fn empty_resolved() -> ResolvedRun {
296        ResolvedRun {
297            agent: BaseAgent::default(),
298            tools: HashMap::new(),
299            run_policy: tirea_contract::RunPolicy::new(),
300            parent_tool_call_id: None,
301        }
302    }
303
304    fn build_read_only_ctx<'a>(phase: Phase, step: &'a StepContext<'a>) -> ReadOnlyContext<'a> {
305        let mut ctx = ReadOnlyContext::new(
306            phase,
307            step.thread_id(),
308            step.messages(),
309            step.run_policy(),
310            step.ctx().doc(),
311        );
312        if let Some(gate) = step.gate.as_ref() {
313            ctx = ctx.with_tool_info(&gate.name, &gate.id, Some(&gate.args));
314            if let Some(result) = gate.result.as_ref() {
315                ctx = ctx.with_tool_result(result);
316            }
317        }
318        if let Some(call_id) = step.tool_call_id() {
319            if let Ok(Some(resume)) = step.ctx().resume_input_for(call_id) {
320                ctx = ctx.with_resume_input(resume);
321            }
322        }
323        ctx
324    }
325
326    struct MarkerPlugin;
327
328    #[async_trait]
329    impl AgentBehavior for MarkerPlugin {
330        fn id(&self) -> &str {
331            "marker_plugin"
332        }
333    }
334
335    #[test]
336    fn injects_frontend_tools_into_resolved() {
337        let request = RunAgentInput {
338            thread_id: "t1".to_string(),
339            run_id: "r1".to_string(),
340            messages: vec![Message::user("hello")],
341            tools: vec![
342                tirea_protocol_ag_ui::Tool {
343                    name: "copyToClipboard".to_string(),
344                    description: "copy".to_string(),
345                    parameters: Some(json!({
346                        "type": "object",
347                        "properties": {
348                            "text": { "type": "string" }
349                        },
350                        "required": ["text"]
351                    })),
352                    execute: ToolExecutionLocation::Frontend,
353                },
354                tirea_protocol_ag_ui::Tool::backend("search", "backend search"),
355            ],
356            context: vec![],
357            state: None,
358            parent_run_id: None,
359            parent_thread_id: None,
360            model: None,
361            system_prompt: None,
362            config: None,
363            forwarded_props: None,
364        };
365
366        let mut resolved = empty_resolved();
367        apply_agui_extensions(&mut resolved, &request);
368        assert_eq!(resolved.agent.tool_executor.name(), "parallel_streaming");
369        assert!(resolved.tools.contains_key("copyToClipboard"));
370        // Only 1 frontend tool (backend tools are not stubs)
371        assert_eq!(resolved.tools.len(), 1);
372        // FrontendToolPendingPlugin behavior
373        assert!(resolved.agent.behavior.id().contains("agui_frontend_tools"));
374    }
375
376    #[test]
377    fn no_behaviors_added_when_only_decisions_are_present() {
378        let request = RunAgentInput {
379            thread_id: "t1".to_string(),
380            run_id: "r1".to_string(),
381            messages: vec![
382                Message::user("hello"),
383                Message::tool("true", "interaction_1"),
384            ],
385            tools: vec![],
386            context: vec![],
387            state: None,
388            parent_run_id: None,
389            parent_thread_id: None,
390            model: None,
391            system_prompt: None,
392            config: None,
393            forwarded_props: None,
394        };
395
396        let mut resolved = empty_resolved();
397        apply_agui_extensions(&mut resolved, &request);
398        assert!(resolved.tools.is_empty());
399        assert_eq!(resolved.agent.behavior.id(), "noop");
400    }
401
402    #[test]
403    fn no_behaviors_added_for_non_boolean_decision_payload_without_frontend_tools() {
404        let request = RunAgentInput {
405            thread_id: "t1".to_string(),
406            run_id: "r1".to_string(),
407            messages: vec![
408                Message::user("hello"),
409                Message::tool(r#"{"todo":"ship starter"}"#, "call_copy_1"),
410            ],
411            tools: vec![],
412            context: vec![],
413            state: Some(json!({
414                "__suspended_tool_calls": {
415                    "calls": {
416                        "call_copy_1": {
417                            "call_id": "call_copy_1",
418                            "tool_name": "copyToClipboard",
419                            "suspension": {
420                                "id": "call_copy_1",
421                                "action": "tool:copyToClipboard"
422                            },
423                            "arguments": {},
424                            "pending": {
425                                "id": "call_copy_1",
426                                "name": "copyToClipboard",
427                                "arguments": {}
428                            },
429                            "resume_mode": "use_decision_as_tool_result"
430                        }
431                    }
432                }
433            })),
434            parent_run_id: None,
435            parent_thread_id: None,
436            model: None,
437            system_prompt: None,
438            config: None,
439            forwarded_props: None,
440        };
441
442        let mut resolved = empty_resolved();
443        apply_agui_extensions(&mut resolved, &request);
444        assert!(resolved.tools.is_empty());
445        assert_eq!(resolved.agent.behavior.id(), "noop");
446    }
447
448    #[test]
449    fn injects_frontend_behavior_even_when_decisions_are_present() {
450        let request = RunAgentInput {
451            thread_id: "t1".to_string(),
452            run_id: "r1".to_string(),
453            messages: vec![Message::user("hello"), Message::tool("true", "call_1")],
454            tools: vec![tirea_protocol_ag_ui::Tool {
455                name: "copyToClipboard".to_string(),
456                description: "copy".to_string(),
457                parameters: None,
458                execute: ToolExecutionLocation::Frontend,
459            }],
460            context: vec![],
461            state: None,
462            parent_run_id: None,
463            parent_thread_id: None,
464            model: None,
465            system_prompt: None,
466            config: None,
467            forwarded_props: None,
468        };
469
470        let mut resolved = empty_resolved();
471        apply_agui_extensions(&mut resolved, &request);
472        assert!(resolved.tools.contains_key("copyToClipboard"));
473        // FrontendToolPendingPlugin behavior only
474        assert_eq!(resolved.agent.behavior.id(), "agui_frontend_tools");
475    }
476
477    #[test]
478    fn composes_frontend_pending_behavior_with_existing_behavior() {
479        let request = RunAgentInput {
480            thread_id: "t1".to_string(),
481            run_id: "r1".to_string(),
482            messages: vec![Message::user("hello")],
483            tools: vec![tirea_protocol_ag_ui::Tool {
484                name: "copyToClipboard".to_string(),
485                description: "copy".to_string(),
486                parameters: None,
487                execute: ToolExecutionLocation::Frontend,
488            }],
489            context: vec![],
490            state: None,
491            parent_run_id: None,
492            parent_thread_id: None,
493            model: None,
494            system_prompt: None,
495            config: None,
496            forwarded_props: None,
497        };
498
499        let mut resolved = empty_resolved();
500        add_behavior_mut(&mut resolved.agent, Arc::new(MarkerPlugin));
501
502        apply_agui_extensions(&mut resolved, &request);
503
504        let behavior_id = resolved.agent.behavior.id();
505        assert!(
506            behavior_id.contains("agui_frontend_tools"),
507            "behavior should contain frontend tools, got: {behavior_id}"
508        );
509        assert!(
510            behavior_id.contains("marker_plugin"),
511            "behavior should contain marker plugin, got: {behavior_id}"
512        );
513    }
514
515    #[test]
516    fn no_changes_without_frontend_or_response_data() {
517        let request = RunAgentInput::new("t1", "r1").with_message(Message::user("hello"));
518        let mut resolved = empty_resolved();
519        apply_agui_extensions(&mut resolved, &request);
520        assert_eq!(resolved.agent.tool_executor.name(), "parallel_streaming");
521        assert!(resolved.tools.is_empty());
522        assert_eq!(resolved.agent.behavior.id(), "noop");
523    }
524
525    #[tokio::test]
526    async fn frontend_pending_plugin_marks_frontend_call_as_pending() {
527        let plugin =
528            FrontendToolPendingPlugin::new(["copyToClipboard".to_string()].into_iter().collect());
529        let fixture = TestFixture::new();
530        let mut step = fixture.step(vec![]);
531        let call = ToolCall::new("call_1", "copyToClipboard", json!({"text":"hello"}));
532        step.gate = Some(ToolGate::from_tool_call(&call));
533
534        let ctx = build_read_only_ctx(Phase::BeforeToolExecute, &step);
535        let actions = plugin.before_tool_execute(&ctx).await;
536        tirea_contract::testing::apply_before_tool_for_test(&mut step, actions);
537
538        let gate = step.gate.as_ref().expect("ToolGate should exist");
539        assert!(gate.pending, "should be pending (suspended)");
540        let ticket = gate
541            .suspend_ticket
542            .as_ref()
543            .expect("should have SuspendTool ticket");
544        assert_eq!(ticket.suspension.action, "tool:copyToClipboard");
545        assert_eq!(ticket.pending.id, "call_1");
546        assert_eq!(ticket.pending.name, "copyToClipboard");
547        assert_eq!(ticket.pending.arguments["text"], "hello");
548        assert_eq!(
549            ticket.resume_mode,
550            tirea_contract::runtime::ToolCallResumeMode::UseDecisionAsToolResult
551        );
552    }
553
554    #[tokio::test]
555    async fn frontend_pending_plugin_resume_sets_tool_result() {
556        let plugin =
557            FrontendToolPendingPlugin::new(["copyToClipboard".to_string()].into_iter().collect());
558        let fixture = TestFixture::new_with_state(json!({
559            "__tool_call_scope": {
560                "call_1": {
561                    "tool_call_state": {
562                        "call_id": "call_1",
563                        "tool_name": "copyToClipboard",
564                        "arguments": { "text": "hello" },
565                        "status": "resuming",
566                        "resume_token": "call_1",
567                        "resume": {
568                            "decision_id": "d1",
569                            "action": "resume",
570                            "result": { "accepted": true },
571                            "updated_at": 1
572                        },
573                        "scratch": null,
574                        "updated_at": 1
575                    }
576                }
577            }
578        }));
579        let mut step = fixture.step(vec![]);
580        let call = ToolCall::new("call_1", "copyToClipboard", json!({"text":"hello"}));
581        step.gate = Some(ToolGate::from_tool_call(&call));
582
583        let ctx = build_read_only_ctx(Phase::BeforeToolExecute, &step);
584        let actions = plugin.before_tool_execute(&ctx).await;
585        tirea_contract::testing::apply_before_tool_for_test(&mut step, actions);
586
587        let gate = step.gate.as_ref().expect("ToolGate should exist");
588        assert!(!gate.blocked, "should not be blocked");
589        assert!(!gate.pending, "should not be pending");
590        let result = gate
591            .result
592            .as_ref()
593            .expect("resume should produce OverrideToolResult");
594        assert_eq!(result.tool_name, "copyToClipboard");
595        assert_eq!(
596            result.status,
597            tirea_contract::runtime::tool_call::ToolStatus::Success
598        );
599        assert_eq!(result.data, json!({"accepted": true}));
600    }
601
602    #[tokio::test]
603    async fn frontend_pending_plugin_cancel_sets_error_result() {
604        let plugin =
605            FrontendToolPendingPlugin::new(["copyToClipboard".to_string()].into_iter().collect());
606        let fixture = TestFixture::new_with_state(json!({
607            "__tool_call_scope": {
608                "call_1": {
609                    "tool_call_state": {
610                        "call_id": "call_1",
611                        "tool_name": "copyToClipboard",
612                        "arguments": { "text": "hello" },
613                        "status": "resuming",
614                        "resume_token": "call_1",
615                        "resume": {
616                            "decision_id": "d1",
617                            "action": "cancel",
618                            "result": null,
619                            "reason": "user denied",
620                            "updated_at": 1
621                        },
622                        "scratch": null,
623                        "updated_at": 1
624                    }
625                }
626            }
627        }));
628        let mut step = fixture.step(vec![]);
629        let call = ToolCall::new("call_1", "copyToClipboard", json!({"text":"hello"}));
630        step.gate = Some(ToolGate::from_tool_call(&call));
631
632        let ctx = build_read_only_ctx(Phase::BeforeToolExecute, &step);
633        let actions = plugin.before_tool_execute(&ctx).await;
634        tirea_contract::testing::apply_before_tool_for_test(&mut step, actions);
635
636        let gate = step.gate.as_ref().expect("ToolGate should exist");
637        assert!(!gate.blocked, "should not be blocked");
638        assert!(!gate.pending, "should not be pending");
639        let result = gate
640            .result
641            .as_ref()
642            .expect("cancel should produce OverrideToolResult");
643        assert_eq!(
644            result.status,
645            tirea_contract::runtime::tool_call::ToolStatus::Error
646        );
647        assert_eq!(result.message.as_deref(), Some("user denied"));
648    }
649
650    #[test]
651    fn injects_context_injection_behavior_when_context_present() {
652        let request = RunAgentInput {
653            thread_id: "t1".to_string(),
654            run_id: "r1".to_string(),
655            messages: vec![Message::user("hello")],
656            tools: vec![],
657            context: vec![Context {
658                description: "Current tasks".to_string(),
659                value: json!(["Review PR", "Write tests"]),
660            }],
661            state: None,
662            parent_run_id: None,
663            parent_thread_id: None,
664            model: None,
665            system_prompt: None,
666            config: None,
667            forwarded_props: None,
668        };
669
670        let mut resolved = empty_resolved();
671        apply_agui_extensions(&mut resolved, &request);
672        assert!(resolved
673            .agent
674            .behavior
675            .id()
676            .contains("agui_context_injection"));
677    }
678
679    #[tokio::test]
680    async fn context_injection_behavior_adds_system_context() {
681        let request = RunAgentInput {
682            thread_id: "t1".to_string(),
683            run_id: "r1".to_string(),
684            messages: vec![Message::user("hello")],
685            tools: vec![],
686            context: vec![Context {
687                description: "Task list".to_string(),
688                value: json!(["Review PR", "Write tests"]),
689            }],
690            state: None,
691            parent_run_id: None,
692            parent_thread_id: None,
693            model: None,
694            system_prompt: None,
695            config: None,
696            forwarded_props: None,
697        };
698
699        let mut resolved = empty_resolved();
700        apply_agui_extensions(&mut resolved, &request);
701        let behavior = &resolved.agent.behavior;
702
703        let fixture = TestFixture::new();
704        let mut step = fixture.step(vec![]);
705        let ctx = build_read_only_ctx(Phase::BeforeInference, &step);
706
707        let actions = behavior.before_inference(&ctx).await;
708        tirea_contract::testing::apply_before_inference_for_test(&mut step, actions);
709
710        assert!(!step.inference.system_context.is_empty());
711        let merged = step.inference.system_context.join("\n");
712        assert!(
713            merged.contains("Task list"),
714            "should contain context description"
715        );
716        assert!(
717            merged.contains("Review PR"),
718            "should contain context values"
719        );
720    }
721
722    #[test]
723    fn no_context_injection_behavior_when_context_empty() {
724        let request = RunAgentInput::new("t1", "r1").with_message(Message::user("hello"));
725        let mut resolved = empty_resolved();
726        apply_agui_extensions(&mut resolved, &request);
727        assert!(!resolved
728            .agent
729            .behavior
730            .id()
731            .contains("agui_context_injection"));
732    }
733
734    #[test]
735    fn applies_request_model_and_system_prompt_overrides() {
736        let request = RunAgentInput::new("t1", "r1")
737            .with_message(Message::user("hello"))
738            .with_model("gpt-4.1")
739            .with_system_prompt("You are precise.");
740
741        let mut resolved = empty_resolved();
742        resolved.agent.model = "base-model".to_string();
743        resolved.agent.system_prompt = "base-prompt".to_string();
744
745        apply_agui_extensions(&mut resolved, &request);
746
747        assert_eq!(resolved.agent.model, "gpt-4.1");
748        assert_eq!(resolved.agent.system_prompt, "You are precise.");
749    }
750
751    #[test]
752    fn ignores_non_runtime_agui_config_fields() {
753        let request = RunAgentInput::new("t1", "r1")
754            .with_message(Message::user("hello"))
755            .with_state(json!({"k":"v"}))
756            .with_forwarded_props(json!({"session":"abc"}));
757        let mut request = request;
758        request.config = Some(json!({"temperature": 0.2}));
759
760        let mut resolved = empty_resolved();
761        apply_agui_extensions(&mut resolved, &request);
762
763        let options = resolved
764            .agent
765            .chat_options
766            .as_ref()
767            .expect("default chat options should be preserved");
768        assert_eq!(options.capture_usage, Some(true));
769        assert_eq!(options.capture_reasoning_content, Some(true));
770    }
771
772    #[test]
773    fn applies_chat_options_overrides_from_agui_config() {
774        let mut request = RunAgentInput::new("t1", "r1").with_message(Message::user("hello"));
775        request.config = Some(json!({
776            "captureReasoningContent": true,
777            "normalizeReasoningContent": true,
778            "reasoningEffort": "high"
779        }));
780
781        let mut resolved = empty_resolved();
782        apply_agui_extensions(&mut resolved, &request);
783
784        let options = resolved
785            .agent
786            .chat_options
787            .expect("chat options should exist");
788        assert_eq!(options.capture_reasoning_content, Some(true));
789        assert_eq!(options.normalize_reasoning_content, Some(true));
790    }
791
792    #[test]
793    fn applies_tool_execution_mode_override_from_agui_config() {
794        let mut request = RunAgentInput::new("t1", "r1").with_message(Message::user("hello"));
795        request.config = Some(json!({
796            "toolExecutionMode": "parallel_batch_approval"
797        }));
798
799        let mut resolved = empty_resolved();
800        apply_agui_extensions(&mut resolved, &request);
801
802        assert_eq!(
803            resolved.agent.tool_executor.name(),
804            "parallel_batch_approval"
805        );
806    }
807
808    #[test]
809    fn applies_nested_chat_options_overrides_from_agui_config() {
810        let mut request = RunAgentInput::new("t1", "r1").with_message(Message::user("hello"));
811        request.config = Some(json!({
812            "chat_options": {
813                "capture_reasoning_content": false,
814                "reasoning_effort": 256
815            }
816        }));
817
818        let mut resolved = empty_resolved();
819        apply_agui_extensions(&mut resolved, &request);
820
821        let options = resolved
822            .agent
823            .chat_options
824            .expect("chat options should exist");
825        assert_eq!(options.capture_reasoning_content, Some(false));
826        assert_eq!(options.normalize_reasoning_content, None);
827    }
828}
tirea_agentos_server/protocol/ag_ui/runtime.rs

tirea_agentos_server/protocol/ag_ui/
runtime.rs
