Skip to main content

context_harness/agent_runtime/
control.rs

1use crate::{
2    agent_store::{BlockedReason, NeedsUserInputReason, RunOutcome},
3    traits::{RunControlKind, Tool, ToolContext, ToolRuntimeDispatch},
4};
5use anyhow::{bail, ensure, Result};
6use async_trait::async_trait;
7use serde::Deserialize;
8use serde_json::{json, Value};
9
10pub(super) const BLOCKED: &str = "run.blocked";
11pub(super) const REQUEST_USER_INPUT: &str = "run.request_user_input";
12
13pub(super) fn is_reserved(name: &str) -> bool {
14    matches!(name, BLOCKED | REQUEST_USER_INPUT)
15}
16
17pub(super) fn register(registry: &mut crate::traits::ToolRegistry) {
18    registry.register(Box::new(ControlTool(RunControlKind::Blocked)));
19    registry.register(Box::new(ControlTool(RunControlKind::RequestUserInput)));
20}
21
22pub(super) fn outcome(kind: RunControlKind, arguments: Value) -> Result<RunOutcome> {
23    match kind {
24        RunControlKind::Blocked => {
25            let args: BlockedArgs = serde_json::from_value(arguments)?;
26            ensure!(
27                !args.message.trim().is_empty() && args.message.len() <= 8192,
28                "invalid blocked message"
29            );
30            Ok(RunOutcome::Blocked {
31                code: args.code,
32                message: args.message,
33            })
34        }
35        RunControlKind::RequestUserInput => {
36            let args: UserInputArgs = serde_json::from_value(arguments)?;
37            ensure!(
38                !args.question.trim().is_empty() && args.question.len() <= 8192,
39                "invalid user-input question"
40            );
41            Ok(RunOutcome::NeedsUserInput {
42                code: args.code,
43                question: args.question,
44            })
45        }
46    }
47}
48
49#[derive(Deserialize)]
50#[serde(deny_unknown_fields)]
51struct BlockedArgs {
52    code: BlockedReason,
53    message: String,
54}
55
56#[derive(Deserialize)]
57#[serde(deny_unknown_fields)]
58struct UserInputArgs {
59    code: NeedsUserInputReason,
60    question: String,
61}
62
63struct ControlTool(RunControlKind);
64
65#[async_trait]
66impl Tool for ControlTool {
67    fn capabilities(&self) -> Option<Vec<crate::agent_resource::Capability>> {
68        Some(Vec::new())
69    }
70
71    fn runtime_dispatch(&self) -> ToolRuntimeDispatch {
72        ToolRuntimeDispatch::RunControl(self.0)
73    }
74
75    fn name(&self) -> &str {
76        match self.0 {
77            RunControlKind::Blocked => BLOCKED,
78            RunControlKind::RequestUserInput => REQUEST_USER_INPUT,
79        }
80    }
81
82    fn description(&self) -> &str {
83        match self.0 {
84            RunControlKind::Blocked => "Suspend because an external condition prevents progress",
85            RunControlKind::RequestUserInput => "Suspend with one concrete question for the user",
86        }
87    }
88
89    fn parameters_schema(&self) -> Value {
90        match self.0 {
91            RunControlKind::Blocked => json!({
92                "type":"object", "additionalProperties":false,
93                "properties":{
94                    "code":{"type":"string","enum":["external_dependency","environment_unavailable","policy_restriction","resource_unavailable"]},
95                    "message":{"type":"string","minLength":1,"maxLength":8192}
96                }, "required":["code","message"]
97            }),
98            RunControlKind::RequestUserInput => json!({
99                "type":"object", "additionalProperties":false,
100                "properties":{
101                    "code":{"type":"string","enum":["decision_required","information_required"]},
102                    "question":{"type":"string","minLength":1,"maxLength":8192}
103                }, "required":["code","question"]
104            }),
105        }
106    }
107
108    fn validate_arguments(&self, params: &Value) -> Result<()> {
109        outcome(self.0, params.clone()).map(|_| ())
110    }
111
112    async fn execute(&self, _params: Value, _ctx: &ToolContext) -> Result<Value> {
113        bail!("run controls are handled by the agent runtime")
114    }
115}