Skip to main content

phaseo_agent/
lib.rs

1#![forbid(unsafe_code)]
2
3use std::error::Error;
4use std::fmt;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::sync::Arc;
7use std::thread;
8use std::time::{Duration, SystemTime, UNIX_EPOCH};
9
10use phaseo::Phaseo;
11use serde::{Deserialize, Serialize};
12use serde_json::{json, Value};
13
14static RUN_SEQUENCE: AtomicU64 = AtomicU64::new(1);
15
16#[derive(Clone, Debug, PartialEq, Eq)]
17pub struct AgentError {
18    message: String,
19}
20
21impl AgentError {
22    pub fn new(message: impl Into<String>) -> Self {
23        Self {
24            message: message.into(),
25        }
26    }
27
28    pub fn message(&self) -> &str {
29        &self.message
30    }
31}
32
33impl fmt::Display for AgentError {
34    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
35        formatter.write_str(&self.message)
36    }
37}
38
39impl Error for AgentError {}
40
41impl From<phaseo::PhaseoError> for AgentError {
42    fn from(error: phaseo::PhaseoError) -> Self {
43        Self::new(error.to_string())
44    }
45}
46
47#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
48pub struct ToolCall {
49    pub id: String,
50    pub name: String,
51    #[serde(default)]
52    pub input: Value,
53}
54
55#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
56pub struct Message {
57    pub role: String,
58    #[serde(default)]
59    pub content: String,
60    #[serde(default, skip_serializing_if = "Vec::is_empty")]
61    pub tool_calls: Vec<ToolCall>,
62    #[serde(default, skip_serializing_if = "Option::is_none")]
63    pub tool_call_id: Option<String>,
64    #[serde(default, skip_serializing_if = "Option::is_none")]
65    pub name: Option<String>,
66    #[serde(default)]
67    pub is_error: bool,
68}
69
70impl Message {
71    pub fn user(content: impl Into<String>) -> Self {
72        Self {
73            role: "user".to_string(),
74            content: content.into(),
75            ..Self::default()
76        }
77    }
78
79    pub fn assistant(content: impl Into<String>) -> Self {
80        Self {
81            role: "assistant".to_string(),
82            content: content.into(),
83            ..Self::default()
84        }
85    }
86}
87
88#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
89pub struct UsageSummary {
90    pub input_tokens: u64,
91    pub output_tokens: u64,
92    pub cached_tokens: u64,
93    pub total_tokens: u64,
94    pub cost: f64,
95}
96
97impl UsageSummary {
98    fn add(&mut self, other: &Self) {
99        self.input_tokens += other.input_tokens;
100        self.output_tokens += other.output_tokens;
101        self.cached_tokens += other.cached_tokens;
102        self.total_tokens += other.total_tokens;
103        self.cost += other.cost;
104    }
105}
106
107#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
108pub struct ModelRequest {
109    pub agent_id: String,
110    pub model: String,
111    pub instructions: String,
112    pub messages: Vec<Message>,
113    pub tools: Vec<ToolSpec>,
114    #[serde(default)]
115    pub context: Value,
116}
117
118#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
119pub struct ModelResponse {
120    pub message: Message,
121    #[serde(default)]
122    pub usage: UsageSummary,
123    pub request_id: Option<String>,
124    pub provider: Option<String>,
125    pub model: Option<String>,
126    pub finish_reason: Option<String>,
127}
128
129pub trait ModelClient {
130    fn generate(&mut self, request: &ModelRequest) -> Result<ModelResponse, AgentError>;
131}
132
133#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
134pub struct ToolSpec {
135    pub id: String,
136    pub description: String,
137    pub parameters: Value,
138}
139
140pub type ToolExecutor =
141    Arc<dyn Fn(Value, &RuntimeContext) -> Result<Value, AgentError> + Send + Sync>;
142
143#[derive(Clone)]
144pub struct Tool {
145    pub id: String,
146    pub description: String,
147    pub parameters: Value,
148    pub execute: Option<ToolExecutor>,
149    pub require_approval: bool,
150}
151
152impl fmt::Debug for Tool {
153    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
154        formatter
155            .debug_struct("Tool")
156            .field("id", &self.id)
157            .field("description", &self.description)
158            .field("parameters", &self.parameters)
159            .field("has_executor", &self.execute.is_some())
160            .field("require_approval", &self.require_approval)
161            .finish()
162    }
163}
164
165impl Tool {
166    pub fn new<F>(
167        id: impl Into<String>,
168        description: impl Into<String>,
169        parameters: Value,
170        execute: F,
171    ) -> Self
172    where
173        F: Fn(Value, &RuntimeContext) -> Result<Value, AgentError> + Send + Sync + 'static,
174    {
175        Self {
176            id: id.into(),
177            description: description.into(),
178            parameters,
179            execute: Some(Arc::new(execute)),
180            require_approval: false,
181        }
182    }
183
184    pub fn external(
185        id: impl Into<String>,
186        description: impl Into<String>,
187        parameters: Value,
188    ) -> Self {
189        Self {
190            id: id.into(),
191            description: description.into(),
192            parameters,
193            execute: None,
194            require_approval: false,
195        }
196    }
197
198    pub fn require_approval(mut self) -> Self {
199        self.require_approval = true;
200        self
201    }
202
203    fn spec(&self) -> ToolSpec {
204        ToolSpec {
205            id: self.id.clone(),
206            description: self.description.clone(),
207            parameters: self.parameters.clone(),
208        }
209    }
210}
211
212pub fn define_tool(tool: Tool) -> Tool {
213    tool
214}
215
216#[derive(Clone, Debug)]
217pub struct RuntimeContext {
218    pub run_id: String,
219    pub agent_id: String,
220    pub step_index: usize,
221    pub context: Value,
222    pub tool_call: ToolCall,
223}
224
225#[derive(Clone, Debug)]
226pub struct HumanReviewContext {
227    pub run_id: String,
228    pub agent_id: String,
229    pub step_index: usize,
230    pub messages: Vec<Message>,
231    pub response: ModelResponse,
232    pub context: Value,
233}
234
235#[derive(Clone, Debug)]
236pub struct HumanReviewRequest {
237    pub reason: String,
238    pub payload: Value,
239}
240
241pub type HumanReviewer =
242    Arc<dyn Fn(&HumanReviewContext) -> Option<HumanReviewRequest> + Send + Sync>;
243
244#[derive(Clone)]
245pub struct AgentDefinition {
246    pub id: String,
247    pub model: String,
248    pub instructions: String,
249    pub tools: Vec<Tool>,
250    pub max_steps: usize,
251    pub max_retries: usize,
252    pub retry_backoff: Duration,
253    pub human_review: Option<HumanReviewer>,
254}
255
256impl fmt::Debug for AgentDefinition {
257    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
258        formatter
259            .debug_struct("AgentDefinition")
260            .field("id", &self.id)
261            .field("model", &self.model)
262            .field("instructions", &self.instructions)
263            .field("tools", &self.tools)
264            .field("max_steps", &self.max_steps)
265            .field("max_retries", &self.max_retries)
266            .field("retry_backoff", &self.retry_backoff)
267            .field("has_human_review", &self.human_review.is_some())
268            .finish()
269    }
270}
271
272impl AgentDefinition {
273    pub fn new(id: impl Into<String>, model: impl Into<String>) -> Self {
274        Self {
275            id: id.into(),
276            model: model.into(),
277            instructions: String::new(),
278            tools: Vec::new(),
279            max_steps: 8,
280            max_retries: 0,
281            retry_backoff: Duration::from_millis(250),
282            human_review: None,
283        }
284    }
285
286    pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
287        self.instructions = instructions.into();
288        self
289    }
290
291    pub fn tool(mut self, tool: Tool) -> Self {
292        self.tools.push(tool);
293        self
294    }
295
296    pub fn max_steps(mut self, max_steps: usize) -> Self {
297        self.max_steps = max_steps.max(1);
298        self
299    }
300
301    pub fn model_retries(mut self, max_retries: usize, backoff: Duration) -> Self {
302        self.max_retries = max_retries;
303        self.retry_backoff = backoff;
304        self
305    }
306
307    pub fn human_review<F>(mut self, reviewer: F) -> Self
308    where
309        F: Fn(&HumanReviewContext) -> Option<HumanReviewRequest> + Send + Sync + 'static,
310    {
311        self.human_review = Some(Arc::new(reviewer));
312        self
313    }
314}
315
316#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
317pub struct PendingToolCall {
318    pub call: ToolCall,
319    pub kind: String,
320    pub reason: String,
321}
322
323#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
324pub struct HumanPause {
325    pub reason: String,
326    pub payload: Value,
327    pub kind: String,
328    pub pending_tool_calls: Vec<PendingToolCall>,
329}
330
331#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
332pub struct RunStep {
333    pub index: usize,
334    pub status: String,
335    pub model_attempts: usize,
336    pub tool_calls: Vec<ToolCall>,
337    pub request_id: Option<String>,
338    pub provider: Option<String>,
339    pub model: Option<String>,
340    pub finish_reason: Option<String>,
341    pub error: Option<String>,
342    pub usage: UsageSummary,
343}
344
345#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
346pub struct RunRecord {
347    pub id: String,
348    pub agent_id: String,
349    pub model: String,
350    pub max_steps: usize,
351    pub status: String,
352    pub input: Value,
353    pub context: Value,
354    pub step_count: usize,
355    pub pause: Option<HumanPause>,
356    pub stop_reason: Option<String>,
357    pub created_at_ms: u64,
358    pub updated_at_ms: u64,
359}
360
361#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
362pub struct RunResult {
363    pub run: RunRecord,
364    pub steps: Vec<RunStep>,
365    pub output: Value,
366    pub messages: Vec<Message>,
367    pub usage: UsageSummary,
368}
369
370#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
371pub struct AgentEvent {
372    pub event_type: String,
373    pub run_id: String,
374    pub agent_id: String,
375    pub timestamp_ms: u64,
376    pub details: Value,
377}
378
379#[derive(Clone, Debug)]
380pub struct RunOptions {
381    pub input: Value,
382    pub context: Value,
383    pub model: Option<String>,
384    pub max_steps: Option<usize>,
385}
386
387impl RunOptions {
388    pub fn new(input: impl Into<Value>) -> Self {
389        Self {
390            input: input.into(),
391            context: Value::Null,
392            model: None,
393            max_steps: None,
394        }
395    }
396}
397
398#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
399pub struct ToolDecision {
400    pub tool_call_id: String,
401    pub reason: Option<String>,
402}
403
404#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
405pub struct ToolOutput {
406    pub tool_call_id: String,
407    pub output: Value,
408}
409
410#[derive(Clone, Debug)]
411pub struct ContinueOptions {
412    pub result: RunResult,
413    pub human_input: Option<String>,
414    pub approvals: Vec<ToolDecision>,
415    pub tool_outputs: Vec<ToolOutput>,
416}
417
418impl ContinueOptions {
419    pub fn new(result: RunResult) -> Self {
420        Self {
421            result,
422            human_input: None,
423            approvals: Vec::new(),
424            tool_outputs: Vec::new(),
425        }
426    }
427}
428
429pub struct Agent {
430    definition: AgentDefinition,
431}
432
433pub fn create_agent(definition: AgentDefinition) -> Agent {
434    Agent { definition }
435}
436
437impl Agent {
438    pub fn run(
439        &self,
440        client: &mut dyn ModelClient,
441        options: RunOptions,
442    ) -> Result<RunResult, AgentError> {
443        self.run_with_events(client, options, None)
444    }
445
446    pub fn run_with_events(
447        &self,
448        client: &mut dyn ModelClient,
449        options: RunOptions,
450        mut on_event: Option<&mut dyn FnMut(&AgentEvent)>,
451    ) -> Result<RunResult, AgentError> {
452        let timestamp = now_ms();
453        let run_id = new_run_id();
454        let model = options
455            .model
456            .unwrap_or_else(|| self.definition.model.clone());
457        let max_steps = options
458            .max_steps
459            .unwrap_or(self.definition.max_steps)
460            .max(1);
461        let input = options.input;
462        let mut result = RunResult {
463            run: RunRecord {
464                id: run_id.clone(),
465                agent_id: self.definition.id.clone(),
466                model,
467                max_steps,
468                status: "running".to_string(),
469                input: input.clone(),
470                context: options.context,
471                step_count: 0,
472                pause: None,
473                stop_reason: None,
474                created_at_ms: timestamp,
475                updated_at_ms: timestamp,
476            },
477            steps: Vec::new(),
478            output: Value::Null,
479            messages: vec![Message::user(value_to_text(&input))],
480            usage: UsageSummary::default(),
481        };
482        emit(&mut on_event, &result, "run.started", Value::Null);
483        self.drive(client, &mut result, &mut on_event)?;
484        Ok(result)
485    }
486
487    pub fn continue_run(
488        &self,
489        client: &mut dyn ModelClient,
490        options: ContinueOptions,
491    ) -> Result<RunResult, AgentError> {
492        self.continue_with_events(client, options, None)
493    }
494
495    pub fn continue_with_events(
496        &self,
497        client: &mut dyn ModelClient,
498        options: ContinueOptions,
499        mut on_event: Option<&mut dyn FnMut(&AgentEvent)>,
500    ) -> Result<RunResult, AgentError> {
501        let ContinueOptions {
502            mut result,
503            human_input,
504            approvals,
505            tool_outputs,
506        } = options;
507        let pause = result
508            .run
509            .pause
510            .clone()
511            .ok_or_else(|| AgentError::new("Agent run is not waiting for human input"))?;
512
513        if pause.kind == "human_review" {
514            let input = human_input
515                .filter(|value| !value.trim().is_empty())
516                .ok_or_else(|| AgentError::new("Human input is required to resume this run"))?;
517            result.messages.push(Message::user(input));
518        } else {
519            self.resolve_pending_tools(&mut result, &pause, &approvals, &tool_outputs)?;
520        }
521
522        result.run.pause = None;
523        result.run.status = "running".to_string();
524        result.run.updated_at_ms = now_ms();
525        emit(&mut on_event, &result, "run.resumed", Value::Null);
526        self.drive(client, &mut result, &mut on_event)?;
527        Ok(result)
528    }
529
530    fn drive(
531        &self,
532        client: &mut dyn ModelClient,
533        result: &mut RunResult,
534        on_event: &mut Option<&mut dyn FnMut(&AgentEvent)>,
535    ) -> Result<(), AgentError> {
536        let max_steps = result.run.max_steps.max(1);
537
538        while result.run.step_count < max_steps {
539            let step_index = result.run.step_count;
540            emit(
541                on_event,
542                result,
543                "model.request.started",
544                json!({"step_index": step_index}),
545            );
546
547            let request = ModelRequest {
548                agent_id: self.definition.id.clone(),
549                model: result.run.model.clone(),
550                instructions: self.definition.instructions.clone(),
551                messages: result.messages.clone(),
552                tools: self.definition.tools.iter().map(Tool::spec).collect(),
553                context: result.run.context.clone(),
554            };
555
556            let (response, attempts) = self.generate_with_retry(client, &request)?;
557            result.run.step_count += 1;
558            result.run.updated_at_ms = now_ms();
559            result.usage.add(&response.usage);
560            result.messages.push(response.message.clone());
561            result.steps.push(RunStep {
562                index: step_index,
563                status: "completed".to_string(),
564                model_attempts: attempts,
565                tool_calls: response.message.tool_calls.clone(),
566                request_id: response.request_id.clone(),
567                provider: response.provider.clone(),
568                model: response.model.clone(),
569                finish_reason: response.finish_reason.clone(),
570                error: None,
571                usage: response.usage.clone(),
572            });
573            emit(
574                on_event,
575                result,
576                "model.response.completed",
577                json!({"step_index": step_index, "attempts": attempts}),
578            );
579
580            if !response.message.tool_calls.is_empty() {
581                let pending = self.pending_tools(&response.message.tool_calls)?;
582                let automatic: Vec<ToolCall> = response
583                    .message
584                    .tool_calls
585                    .iter()
586                    .filter(|call| !pending.iter().any(|item| item.call.id == call.id))
587                    .cloned()
588                    .collect();
589                self.execute_tools(result, &automatic, on_event)?;
590                if !pending.is_empty() {
591                    let pause_kind = if pending.iter().all(|call| call.kind == "external_output") {
592                        "external_output"
593                    } else if pending.iter().all(|call| call.kind == "approval") {
594                        "tool_approval"
595                    } else {
596                        "tool_input"
597                    };
598                    result.run.status = "waiting_for_human".to_string();
599                    result.run.pause = Some(HumanPause {
600                        reason: "Pending tool calls require input".to_string(),
601                        payload: json!({"tool_calls": pending}),
602                        kind: pause_kind.to_string(),
603                        pending_tool_calls: pending,
604                    });
605                    emit(on_event, result, "run.paused", Value::Null);
606                    return Ok(());
607                }
608                continue;
609            }
610
611            if let Some(reviewer) = &self.definition.human_review {
612                let review_context = HumanReviewContext {
613                    run_id: result.run.id.clone(),
614                    agent_id: result.run.agent_id.clone(),
615                    step_index,
616                    messages: result.messages.clone(),
617                    response: response.clone(),
618                    context: result.run.context.clone(),
619                };
620                if let Some(review) = reviewer(&review_context) {
621                    result.run.status = "waiting_for_human".to_string();
622                    result.run.pause = Some(HumanPause {
623                        reason: review.reason,
624                        payload: review.payload,
625                        kind: "human_review".to_string(),
626                        pending_tool_calls: Vec::new(),
627                    });
628                    emit(on_event, result, "run.paused", Value::Null);
629                    return Ok(());
630                }
631            }
632
633            result.output = Value::String(response.message.content);
634            result.run.status = "completed".to_string();
635            emit(on_event, result, "run.completed", Value::Null);
636            return Ok(());
637        }
638
639        result.run.status = "stopped".to_string();
640        result.run.stop_reason = Some(format!("max_steps:{max_steps}"));
641        emit(
642            on_event,
643            result,
644            "run.stopped",
645            json!({"reason": result.run.stop_reason}),
646        );
647        Ok(())
648    }
649
650    fn generate_with_retry(
651        &self,
652        client: &mut dyn ModelClient,
653        request: &ModelRequest,
654    ) -> Result<(ModelResponse, usize), AgentError> {
655        let mut attempts = 0;
656        loop {
657            attempts += 1;
658            match client.generate(request) {
659                Ok(response) => return Ok((response, attempts)),
660                Err(error) if attempts <= self.definition.max_retries => {
661                    if !self.definition.retry_backoff.is_zero() {
662                        thread::sleep(self.definition.retry_backoff);
663                    }
664                    let _ = error;
665                }
666                Err(error) => return Err(error),
667            }
668        }
669    }
670
671    fn pending_tools(&self, calls: &[ToolCall]) -> Result<Vec<PendingToolCall>, AgentError> {
672        calls
673            .iter()
674            .filter_map(|call| {
675                let tool = self
676                    .definition
677                    .tools
678                    .iter()
679                    .find(|tool| tool.id == call.name);
680                match tool {
681                    None => Some(Err(AgentError::new(format!(
682                        "Model requested unknown tool: {}",
683                        call.name
684                    )))),
685                    Some(tool) if tool.require_approval => Some(Ok(PendingToolCall {
686                        call: call.clone(),
687                        kind: "approval".to_string(),
688                        reason: "Tool requires approval".to_string(),
689                    })),
690                    Some(tool) if tool.execute.is_none() => Some(Ok(PendingToolCall {
691                        call: call.clone(),
692                        kind: "external_output".to_string(),
693                        reason: "Tool output must be supplied externally".to_string(),
694                    })),
695                    Some(_) => None,
696                }
697            })
698            .collect()
699    }
700
701    fn execute_tools(
702        &self,
703        result: &mut RunResult,
704        calls: &[ToolCall],
705        on_event: &mut Option<&mut dyn FnMut(&AgentEvent)>,
706    ) -> Result<(), AgentError> {
707        for call in calls {
708            let tool = self
709                .definition
710                .tools
711                .iter()
712                .find(|tool| tool.id == call.name)
713                .ok_or_else(|| AgentError::new(format!("Unknown tool: {}", call.name)))?;
714            let executor = tool.execute.as_ref().ok_or_else(|| {
715                AgentError::new(format!("Tool requires external output: {}", call.name))
716            })?;
717            let runtime = RuntimeContext {
718                run_id: result.run.id.clone(),
719                agent_id: result.run.agent_id.clone(),
720                step_index: result.run.step_count.saturating_sub(1),
721                context: result.run.context.clone(),
722                tool_call: call.clone(),
723            };
724            emit(
725                on_event,
726                result,
727                "tool.started",
728                json!({"tool_call_id": call.id, "name": call.name}),
729            );
730            let (output, is_error) = match executor(call.input.clone(), &runtime) {
731                Ok(output) => (output, false),
732                Err(error) => (json!({"error": error.to_string()}), true),
733            };
734            result.messages.push(Message {
735                role: "tool".to_string(),
736                content: value_to_text(&output),
737                tool_call_id: Some(call.id.clone()),
738                name: Some(call.name.clone()),
739                is_error,
740                ..Message::default()
741            });
742            emit(
743                on_event,
744                result,
745                "tool.completed",
746                json!({"tool_call_id": call.id, "name": call.name, "is_error": is_error}),
747            );
748        }
749        Ok(())
750    }
751
752    fn resolve_pending_tools(
753        &self,
754        result: &mut RunResult,
755        pause: &HumanPause,
756        approvals: &[ToolDecision],
757        tool_outputs: &[ToolOutput],
758    ) -> Result<(), AgentError> {
759        for pending in &pause.pending_tool_calls {
760            match pending.kind.as_str() {
761                "external_output" => {
762                    if let Some(output) = tool_outputs
763                        .iter()
764                        .find(|output| output.tool_call_id == pending.call.id)
765                    {
766                        result.messages.push(Message {
767                            role: "tool".to_string(),
768                            content: value_to_text(&output.output),
769                            tool_call_id: Some(pending.call.id.clone()),
770                            name: Some(pending.call.name.clone()),
771                            ..Message::default()
772                        });
773                        continue;
774                    }
775                }
776                "approval" => {
777                    let approved = approvals
778                        .iter()
779                        .any(|decision| decision.tool_call_id == pending.call.id);
780                    if approved {
781                        let tool = self
782                            .definition
783                            .tools
784                            .iter()
785                            .find(|tool| tool.id == pending.call.name)
786                            .ok_or_else(|| {
787                                AgentError::new(format!(
788                                    "Model requested unknown tool: {}",
789                                    pending.call.name
790                                ))
791                            })?;
792                        if tool.execute.is_some() {
793                            self.execute_tools(
794                                result,
795                                std::slice::from_ref(&pending.call),
796                                &mut None,
797                            )?;
798                            continue;
799                        }
800                        if let Some(output) = tool_outputs
801                            .iter()
802                            .find(|output| output.tool_call_id == pending.call.id)
803                        {
804                            result.messages.push(Message {
805                                role: "tool".to_string(),
806                                content: value_to_text(&output.output),
807                                tool_call_id: Some(pending.call.id.clone()),
808                                name: Some(pending.call.name.clone()),
809                                ..Message::default()
810                            });
811                            continue;
812                        }
813                    }
814                }
815                kind => {
816                    return Err(AgentError::new(format!(
817                        "Unknown pending tool call kind {kind} for {}",
818                        pending.call.id
819                    )));
820                }
821            }
822
823            return Err(AgentError::new(format!(
824                "No approval or external output supplied for tool call {}",
825                pending.call.id
826            )));
827        }
828        Ok(())
829    }
830}
831
832pub struct GatewayAgentClient {
833    client: Phaseo,
834    model: String,
835}
836
837impl GatewayAgentClient {
838    pub fn new(client: Phaseo, model: impl Into<String>) -> Self {
839        Self {
840            client,
841            model: model.into(),
842        }
843    }
844
845    pub fn from_env(model: impl Into<String>) -> Result<Self, AgentError> {
846        Ok(Self::new(
847            Phaseo::from_env()?
848                .with_header("X-Phaseo-Client", "phaseo-agent-rust")
849                .with_header("X-Phaseo-Client-Version", env!("CARGO_PKG_VERSION")),
850            model,
851        ))
852    }
853}
854
855pub fn create_gateway_agent_client(
856    model: impl Into<String>,
857) -> Result<GatewayAgentClient, AgentError> {
858    GatewayAgentClient::from_env(model)
859}
860
861impl ModelClient for GatewayAgentClient {
862    fn generate(&mut self, request: &ModelRequest) -> Result<ModelResponse, AgentError> {
863        let model = if request.model.trim().is_empty() {
864            &self.model
865        } else {
866            &request.model
867        };
868        let input = response_input(&request.messages);
869        let tools: Vec<Value> = request
870            .tools
871            .iter()
872            .map(|tool| {
873                json!({
874                    "type": "function",
875                    "name": tool.id,
876                    "description": tool.description,
877                    "parameters": tool.parameters,
878                })
879            })
880            .collect();
881        let mut body = json!({
882            "model": model,
883            "input": input,
884            "metadata": {"phaseo_agent_id": request.agent_id},
885        });
886        if !request.instructions.is_empty() {
887            body["instructions"] = Value::String(request.instructions.clone());
888        }
889        if !tools.is_empty() {
890            body["tools"] = Value::Array(tools);
891        }
892
893        let response = self.client.responses(&body)?;
894        model_response_from_body(response.body, response.request_id)
895    }
896}
897
898fn response_input(messages: &[Message]) -> Vec<Value> {
899    let mut input = Vec::new();
900    for message in messages {
901        if message.role == "tool" {
902            input.push(json!({
903                "type": "function_call_output",
904                "call_id": message.tool_call_id,
905                "output": message.content,
906            }));
907            continue;
908        }
909        if !message.content.is_empty() {
910            input.push(json!({
911                "type": "message",
912                "role": message.role,
913                "content": message.content,
914            }));
915        }
916        for call in &message.tool_calls {
917            input.push(json!({
918                "type": "function_call",
919                "call_id": call.id,
920                "name": call.name,
921                "arguments": call.input.to_string(),
922            }));
923        }
924    }
925    input
926}
927
928fn model_response_from_body(
929    body: Value,
930    header_request_id: Option<String>,
931) -> Result<ModelResponse, AgentError> {
932    let mut text = body
933        .get("output_text")
934        .and_then(Value::as_str)
935        .unwrap_or_default()
936        .to_string();
937    let mut tool_calls = Vec::new();
938    if let Some(output) = body.get("output").and_then(Value::as_array) {
939        for item in output {
940            match item.get("type").and_then(Value::as_str) {
941                Some("function_call") => {
942                    let raw = item
943                        .get("arguments")
944                        .and_then(Value::as_str)
945                        .unwrap_or("{}");
946                    tool_calls.push(ToolCall {
947                        id: item
948                            .get("call_id")
949                            .or_else(|| item.get("id"))
950                            .and_then(Value::as_str)
951                            .unwrap_or_default()
952                            .to_string(),
953                        name: item
954                            .get("name")
955                            .and_then(Value::as_str)
956                            .unwrap_or_default()
957                            .to_string(),
958                        input: serde_json::from_str(raw)
959                            .unwrap_or_else(|_| Value::String(raw.into())),
960                    });
961                }
962                Some("message") if text.is_empty() => {
963                    text = item
964                        .get("content")
965                        .and_then(Value::as_array)
966                        .into_iter()
967                        .flatten()
968                        .filter_map(|part| part.get("text").and_then(Value::as_str))
969                        .collect::<Vec<_>>()
970                        .join("");
971                }
972                _ => {}
973            }
974        }
975    }
976
977    if text.is_empty() && tool_calls.is_empty() {
978        return Err(AgentError::new(
979            "Phaseo response contained neither output text nor tool calls",
980        ));
981    }
982
983    let usage = body.get("usage").cloned().unwrap_or(Value::Null);
984    let input_tokens = integer(&usage, &["input_tokens", "prompt_tokens"]);
985    let output_tokens = integer(&usage, &["output_tokens", "completion_tokens"]);
986    let total_tokens = integer(&usage, &["total_tokens"]).max(input_tokens + output_tokens);
987    Ok(ModelResponse {
988        message: Message {
989            role: "assistant".to_string(),
990            content: text,
991            tool_calls,
992            ..Message::default()
993        },
994        usage: UsageSummary {
995            input_tokens,
996            output_tokens,
997            cached_tokens: usage
998                .pointer("/input_tokens_details/cached_tokens")
999                .and_then(Value::as_u64)
1000                .unwrap_or_else(|| integer(&usage, &["cached_tokens", "cache_read_input_tokens"])),
1001            total_tokens,
1002            cost: body
1003                .pointer("/meta/cost")
1004                .or_else(|| body.get("cost"))
1005                .and_then(Value::as_f64)
1006                .unwrap_or_default(),
1007        },
1008        request_id: header_request_id
1009            .or_else(|| body.get("id").and_then(Value::as_str).map(str::to_string)),
1010        provider: body
1011            .get("provider")
1012            .and_then(Value::as_str)
1013            .map(str::to_string),
1014        model: body
1015            .get("model")
1016            .and_then(Value::as_str)
1017            .map(str::to_string),
1018        finish_reason: body
1019            .get("status")
1020            .and_then(Value::as_str)
1021            .map(str::to_string),
1022    })
1023}
1024
1025fn integer(value: &Value, keys: &[&str]) -> u64 {
1026    keys.iter()
1027        .find_map(|key| value.get(key).and_then(Value::as_u64))
1028        .unwrap_or_default()
1029}
1030
1031fn value_to_text(value: &Value) -> String {
1032    value
1033        .as_str()
1034        .map(str::to_string)
1035        .unwrap_or_else(|| value.to_string())
1036}
1037
1038fn now_ms() -> u64 {
1039    SystemTime::now()
1040        .duration_since(UNIX_EPOCH)
1041        .unwrap_or_default()
1042        .as_millis()
1043        .try_into()
1044        .unwrap_or(u64::MAX)
1045}
1046
1047fn new_run_id() -> String {
1048    let sequence = RUN_SEQUENCE.fetch_add(1, Ordering::Relaxed);
1049    format!("run_{}_{sequence}", now_ms())
1050}
1051
1052fn emit(
1053    handler: &mut Option<&mut dyn FnMut(&AgentEvent)>,
1054    result: &RunResult,
1055    event_type: &str,
1056    details: Value,
1057) {
1058    if let Some(handler) = handler.as_deref_mut() {
1059        handler(&AgentEvent {
1060            event_type: event_type.to_string(),
1061            run_id: result.run.id.clone(),
1062            agent_id: result.run.agent_id.clone(),
1063            timestamp_ms: now_ms(),
1064            details,
1065        });
1066    }
1067}
1068
1069#[cfg(test)]
1070mod tests {
1071    use super::*;
1072
1073    struct ToolLoopClient {
1074        calls: usize,
1075    }
1076
1077    impl ModelClient for ToolLoopClient {
1078        fn generate(&mut self, _request: &ModelRequest) -> Result<ModelResponse, AgentError> {
1079            self.calls += 1;
1080            if self.calls == 1 {
1081                return Ok(ModelResponse {
1082                    message: Message {
1083                        role: "assistant".to_string(),
1084                        tool_calls: vec![ToolCall {
1085                            id: "call_1".to_string(),
1086                            name: "lookup".to_string(),
1087                            input: json!({"slug": "presets"}),
1088                        }],
1089                        ..Message::default()
1090                    },
1091                    ..ModelResponse::default()
1092                });
1093            }
1094            Ok(ModelResponse {
1095                message: Message::assistant("Presets define stable routing defaults."),
1096                ..ModelResponse::default()
1097            })
1098        }
1099    }
1100
1101    #[test]
1102    fn executes_a_tool_loop_and_emits_events() {
1103        let tool = Tool::new(
1104            "lookup",
1105            "Lookup documentation",
1106            json!({"type": "object"}),
1107            |input, _context| Ok(json!({"slug": input["slug"], "ok": true})),
1108        );
1109        let agent = create_agent(
1110            AgentDefinition::new("support", "openai/gpt-5.4-nano")
1111                .instructions("Use tools when helpful")
1112                .tool(tool),
1113        );
1114        let mut client = ToolLoopClient { calls: 0 };
1115        let mut events = Vec::new();
1116        let mut capture = |event: &AgentEvent| events.push(event.event_type.clone());
1117        let result = agent
1118            .run_with_events(
1119                &mut client,
1120                RunOptions::new("Explain presets"),
1121                Some(&mut capture),
1122            )
1123            .unwrap();
1124
1125        assert_eq!(result.run.status, "completed");
1126        assert_eq!(
1127            result.output,
1128            json!("Presets define stable routing defaults.")
1129        );
1130        assert_eq!(result.steps.len(), 2);
1131        assert!(result.messages.iter().any(|message| message.role == "tool"));
1132        assert!(events.iter().any(|event| event == "tool.completed"));
1133    }
1134
1135    struct RetryReviewClient {
1136        calls: usize,
1137        models: Vec<String>,
1138    }
1139
1140    impl ModelClient for RetryReviewClient {
1141        fn generate(&mut self, request: &ModelRequest) -> Result<ModelResponse, AgentError> {
1142            self.calls += 1;
1143            self.models.push(request.model.clone());
1144            if self.calls == 1 {
1145                return Err(AgentError::new("temporary failure"));
1146            }
1147            Ok(ModelResponse {
1148                message: Message::assistant("Deploy the change"),
1149                ..ModelResponse::default()
1150            })
1151        }
1152    }
1153
1154    #[test]
1155    fn retries_pauses_and_resumes_human_review() {
1156        let agent = create_agent(
1157            AgentDefinition::new("review", "openai/gpt-5.4-nano")
1158                .model_retries(1, Duration::ZERO)
1159                .human_review(|context| {
1160                    if context
1161                        .messages
1162                        .iter()
1163                        .any(|message| message.role == "user" && message.content == "approved")
1164                    {
1165                        None
1166                    } else {
1167                        Some(HumanReviewRequest {
1168                            reason: "Approve deployment".to_string(),
1169                            payload: json!({"output": context.response.message.content}),
1170                        })
1171                    }
1172                }),
1173        );
1174        let mut client = RetryReviewClient {
1175            calls: 0,
1176            models: Vec::new(),
1177        };
1178        let mut run_options = RunOptions::new("Prepare deployment");
1179        run_options.model = Some("openai/override-model".to_string());
1180        run_options.max_steps = Some(3);
1181        let paused = agent.run(&mut client, run_options).unwrap();
1182        assert_eq!(paused.run.status, "waiting_for_human");
1183        assert_eq!(paused.steps[0].model_attempts, 2);
1184        assert_eq!(paused.run.model, "openai/override-model");
1185        assert_eq!(paused.run.max_steps, 3);
1186
1187        let mut options = ContinueOptions::new(paused);
1188        options.human_input = Some("approved".to_string());
1189        let completed = agent.continue_run(&mut client, options).unwrap();
1190        assert_eq!(completed.run.status, "completed");
1191        assert_eq!(completed.output, json!("Deploy the change"));
1192        assert!(client
1193            .models
1194            .iter()
1195            .all(|model| model == "openai/override-model"));
1196        serde_json::to_string(&completed).unwrap();
1197    }
1198
1199    struct MixedToolClient {
1200        calls: usize,
1201    }
1202
1203    impl ModelClient for MixedToolClient {
1204        fn generate(&mut self, _request: &ModelRequest) -> Result<ModelResponse, AgentError> {
1205            self.calls += 1;
1206            if self.calls == 1 {
1207                return Ok(ModelResponse {
1208                    message: Message {
1209                        role: "assistant".to_string(),
1210                        tool_calls: vec![
1211                            ToolCall {
1212                                id: "call_local".to_string(),
1213                                name: "local".to_string(),
1214                                input: json!({"value": 1}),
1215                            },
1216                            ToolCall {
1217                                id: "call_external".to_string(),
1218                                name: "external".to_string(),
1219                                input: json!({"value": 2}),
1220                            },
1221                        ],
1222                        ..Message::default()
1223                    },
1224                    ..ModelResponse::default()
1225                });
1226            }
1227            Ok(ModelResponse {
1228                message: Message::assistant("Both tool results received."),
1229                ..ModelResponse::default()
1230            })
1231        }
1232    }
1233
1234    #[test]
1235    fn executes_automatic_tools_before_pausing_for_external_outputs() {
1236        let agent = create_agent(
1237            AgentDefinition::new("mixed-tools", "openai/gpt-5.4-nano")
1238                .tool(Tool::new(
1239                    "local",
1240                    "Run locally",
1241                    json!({"type": "object"}),
1242                    |input, _context| Ok(json!({"local": input["value"]})),
1243                ))
1244                .tool(Tool::external(
1245                    "external",
1246                    "Run externally",
1247                    json!({"type": "object"}),
1248                )),
1249        );
1250        let mut client = MixedToolClient { calls: 0 };
1251        let paused = agent
1252            .run(&mut client, RunOptions::new("Use both tools"))
1253            .unwrap();
1254
1255        let pause = paused.run.pause.as_ref().unwrap();
1256        assert_eq!(pause.kind, "external_output");
1257        assert_eq!(pause.pending_tool_calls.len(), 1);
1258        assert_eq!(pause.pending_tool_calls[0].call.id, "call_external");
1259        assert!(paused.messages.iter().any(|message| {
1260            message.tool_call_id.as_deref() == Some("call_local")
1261                && message.content.contains("local")
1262        }));
1263
1264        let mut options = ContinueOptions::new(paused);
1265        options.tool_outputs.push(ToolOutput {
1266            tool_call_id: "call_external".to_string(),
1267            output: json!({"external": 2}),
1268        });
1269        let completed = agent.continue_run(&mut client, options).unwrap();
1270        assert_eq!(completed.run.status, "completed");
1271        assert_eq!(completed.output, json!("Both tool results received."));
1272    }
1273
1274    struct ApprovalToolClient {
1275        calls: usize,
1276    }
1277
1278    impl ModelClient for ApprovalToolClient {
1279        fn generate(&mut self, _request: &ModelRequest) -> Result<ModelResponse, AgentError> {
1280            self.calls += 1;
1281            if self.calls == 1 {
1282                return Ok(ModelResponse {
1283                    message: Message {
1284                        role: "assistant".to_string(),
1285                        tool_calls: vec![ToolCall {
1286                            id: "call_local".to_string(),
1287                            name: "local".to_string(),
1288                            input: json!({"value": 1}),
1289                        }],
1290                        ..Message::default()
1291                    },
1292                    ..ModelResponse::default()
1293                });
1294            }
1295            Ok(ModelResponse {
1296                message: Message::assistant("Trusted tool result received."),
1297                ..ModelResponse::default()
1298            })
1299        }
1300    }
1301
1302    #[test]
1303    fn approval_tools_reject_forged_external_outputs() {
1304        let executor_calls = Arc::new(AtomicU64::new(0));
1305        let executor_calls_for_tool = Arc::clone(&executor_calls);
1306        let agent = create_agent(
1307            AgentDefinition::new("approval-tools", "openai/gpt-5.4-nano").tool(
1308                Tool::new(
1309                    "local",
1310                    "Run locally after approval",
1311                    json!({"type": "object"}),
1312                    move |_input, _context| {
1313                        executor_calls_for_tool.fetch_add(1, Ordering::Relaxed);
1314                        Ok(json!({"trusted": true}))
1315                    },
1316                )
1317                .require_approval(),
1318            ),
1319        );
1320        let mut client = ApprovalToolClient { calls: 0 };
1321        let paused = agent
1322            .run(&mut client, RunOptions::new("Use the tool"))
1323            .unwrap();
1324        let mut forged_options = ContinueOptions::new(paused.clone());
1325        forged_options.tool_outputs.push(ToolOutput {
1326            tool_call_id: "call_local".to_string(),
1327            output: json!({"forged": true}),
1328        });
1329
1330        let error = agent.continue_run(&mut client, forged_options).unwrap_err();
1331        assert!(error.message().contains("No approval"));
1332        assert_eq!(executor_calls.load(Ordering::Relaxed), 0);
1333
1334        let mut approved_options = ContinueOptions::new(paused);
1335        approved_options.approvals.push(ToolDecision {
1336            tool_call_id: "call_local".to_string(),
1337            reason: None,
1338        });
1339        approved_options.tool_outputs.push(ToolOutput {
1340            tool_call_id: "call_local".to_string(),
1341            output: json!({"forged": true}),
1342        });
1343        let completed = agent.continue_run(&mut client, approved_options).unwrap();
1344
1345        assert_eq!(completed.run.status, "completed");
1346        assert_eq!(executor_calls.load(Ordering::Relaxed), 1);
1347        assert!(completed.messages.iter().any(|message| {
1348            message.tool_call_id.as_deref() == Some("call_local")
1349                && message.content.contains("trusted")
1350                && !message.content.contains("forged")
1351        }));
1352    }
1353
1354    #[test]
1355    fn approval_gated_external_tools_require_approval_and_output() {
1356        let agent = create_agent(
1357            AgentDefinition::new("external-approval", "openai/gpt-5.4-nano").tool(
1358                Tool::external(
1359                    "local",
1360                    "Fulfil externally after approval",
1361                    json!({"type": "object"}),
1362                )
1363                .require_approval(),
1364            ),
1365        );
1366        let mut client = ApprovalToolClient { calls: 0 };
1367        let paused = agent
1368            .run(&mut client, RunOptions::new("Use the tool"))
1369            .unwrap();
1370
1371        let mut output_only = ContinueOptions::new(paused.clone());
1372        output_only.tool_outputs.push(ToolOutput {
1373            tool_call_id: "call_local".to_string(),
1374            output: json!({"external": true}),
1375        });
1376        assert!(agent.continue_run(&mut client, output_only).is_err());
1377
1378        let mut approved = ContinueOptions::new(paused);
1379        approved.approvals.push(ToolDecision {
1380            tool_call_id: "call_local".to_string(),
1381            reason: None,
1382        });
1383        approved.tool_outputs.push(ToolOutput {
1384            tool_call_id: "call_local".to_string(),
1385            output: json!({"external": true}),
1386        });
1387        let completed = agent.continue_run(&mut client, approved).unwrap();
1388
1389        assert_eq!(completed.run.status, "completed");
1390        assert!(completed.messages.iter().any(|message| {
1391            message.tool_call_id.as_deref() == Some("call_local")
1392                && message.content.contains("external")
1393        }));
1394    }
1395
1396    #[test]
1397    fn parses_gateway_tool_calls() {
1398        let response = model_response_from_body(
1399            json!({
1400                "id": "resp_1",
1401                "model": "test/model",
1402                "status": "completed",
1403                "output": [{
1404                    "type": "function_call",
1405                    "call_id": "call_1",
1406                    "name": "lookup",
1407                    "arguments": "{\"slug\":\"rust\"}"
1408                }],
1409                "usage": {"input_tokens": 3, "output_tokens": 2}
1410            }),
1411            None,
1412        )
1413        .unwrap();
1414        assert_eq!(response.message.tool_calls[0].input["slug"], "rust");
1415        assert_eq!(response.usage.total_tokens, 5);
1416    }
1417
1418    #[test]
1419    fn reads_nested_responses_cached_token_usage() {
1420        let response = model_response_from_body(
1421            json!({
1422                "id": "resp_cached",
1423                "output_text": "cached",
1424                "usage": {
1425                    "input_tokens": 10,
1426                    "output_tokens": 2,
1427                    "input_tokens_details": {"cached_tokens": 7}
1428                }
1429            }),
1430            None,
1431        )
1432        .unwrap();
1433        assert_eq!(response.usage.cached_tokens, 7);
1434    }
1435}