Skip to main content

swink_agent/agent/
mutation.rs

1use std::sync::Arc;
2
3use crate::loop_::AgentEvent;
4use crate::stream::StreamFn;
5use crate::tool::{AgentTool, ApprovalMode};
6use crate::types::{AgentMessage, ModelSpec, ThinkingLevel};
7
8use super::{Agent, DEFAULT_PLAN_MODE_ADDENDUM};
9
10impl Agent {
11    /// Set the system prompt.
12    pub fn set_system_prompt(&mut self, prompt: impl Into<String>) {
13        self.state.system_prompt = prompt.into();
14    }
15
16    /// Set the model specification, swapping the stream function if a matching
17    /// model was registered via [`with_available_models`](crate::AgentOptions::with_available_models).
18    ///
19    /// If the model actually changes, a [`AgentEvent::ModelCycled`] event is
20    /// dispatched to all subscribers.
21    pub fn set_model(&mut self, model: ModelSpec) {
22        let old = self.state.model.clone();
23        if let Some((_, stream_fn)) = self.model_stream_fns.iter().find(|(m, _)| *m == model) {
24            self.stream_fn = Arc::clone(stream_fn);
25        }
26        self.state.model = model.clone();
27        if old != model {
28            let event = AgentEvent::ModelCycled {
29                old,
30                new: model,
31                reason: "set_model".to_string(),
32            };
33            self.dispatch_event(&event);
34        }
35    }
36
37    /// Set the model specification and stream function explicitly, bypassing
38    /// the [`available_models`](crate::AgentState::available_models) lookup.
39    pub fn set_model_with_stream(&mut self, model: ModelSpec, stream_fn: Arc<dyn StreamFn>) {
40        let old = self.state.model.clone();
41        self.state.model = model.clone();
42        self.stream_fn = stream_fn;
43        if old != model {
44            let event = AgentEvent::ModelCycled {
45                old,
46                new: model,
47                reason: "set_model_with_stream".to_string(),
48            };
49            self.dispatch_event(&event);
50        }
51    }
52
53    /// Set the thinking level on the current model.
54    pub const fn set_thinking_level(&mut self, level: ThinkingLevel) {
55        self.state.model.thinking_level = level;
56    }
57
58    /// Replace the tool set.
59    pub fn set_tools(&mut self, tools: Vec<Arc<dyn AgentTool>>) {
60        super::assert_unique_tool_names(&tools);
61        self.state.tools = tools;
62    }
63
64    /// Add a tool, replacing any existing tool with the same name.
65    pub fn add_tool(&mut self, tool: Arc<dyn AgentTool>) {
66        let name = tool.name();
67        self.state.tools.retain(|t| t.name() != name);
68        self.state.tools.push(tool);
69    }
70
71    /// Remove a tool by name. Returns `true` if a tool was found and removed.
72    pub fn remove_tool(&mut self, name: &str) -> bool {
73        let before = self.state.tools.len();
74        self.state.tools.retain(|t| t.name() != name);
75        self.state.tools.len() < before
76    }
77
78    /// Return the current approval mode.
79    pub const fn approval_mode(&self) -> ApprovalMode {
80        self.approval_mode
81    }
82
83    /// Set the approval mode at runtime.
84    pub const fn set_approval_mode(&mut self, mode: ApprovalMode) {
85        self.approval_mode = mode;
86    }
87
88    /// Find a tool by name.
89    #[must_use]
90    pub fn find_tool(&self, name: &str) -> Option<&Arc<dyn AgentTool>> {
91        self.state.tools.iter().find(|t| t.name() == name)
92    }
93
94    /// Return tools matching a predicate.
95    #[must_use]
96    pub fn tools_matching(
97        &self,
98        predicate: impl Fn(&dyn AgentTool) -> bool,
99    ) -> Vec<&Arc<dyn AgentTool>> {
100        self.state
101            .tools
102            .iter()
103            .filter(|t| predicate(t.as_ref()))
104            .collect()
105    }
106
107    /// Return tools belonging to the given namespace.
108    #[must_use]
109    pub fn tools_in_namespace(&self, namespace: &str) -> Vec<&Arc<dyn AgentTool>> {
110        self.state
111            .tools
112            .iter()
113            .filter(|t| {
114                t.metadata()
115                    .and_then(|m| m.namespace)
116                    .is_some_and(|ns| ns == namespace)
117            })
118            .collect()
119    }
120
121    /// Replace the entire message history.
122    pub fn set_messages(&mut self, messages: Vec<AgentMessage>) {
123        self.state.messages = messages;
124    }
125
126    /// Append messages to the history.
127    pub fn append_messages(&mut self, messages: Vec<AgentMessage>) {
128        self.state.messages.extend(messages);
129    }
130
131    /// Clear the message history.
132    pub fn clear_messages(&mut self) {
133        self.state.messages.clear();
134    }
135
136    /// Enter plan mode: restrict to read-only tools and append plan instructions.
137    pub fn enter_plan_mode(&mut self) -> (Vec<Arc<dyn AgentTool>>, String) {
138        let state = &mut self.state;
139        let saved_tools = state.tools.clone();
140        let saved_prompt = state.system_prompt.clone();
141
142        let read_only: Vec<Arc<dyn AgentTool>> = saved_tools
143            .iter()
144            .filter(|tool| !tool.requires_approval())
145            .cloned()
146            .collect();
147        state.tools = read_only;
148
149        let addendum = self
150            .plan_mode_addendum
151            .as_deref()
152            .unwrap_or(DEFAULT_PLAN_MODE_ADDENDUM);
153        state.system_prompt = format!("{}{addendum}", state.system_prompt);
154
155        (saved_tools, saved_prompt)
156    }
157
158    /// Exit plan mode: restore previously saved tools and system prompt.
159    pub fn exit_plan_mode(&mut self, saved_tools: Vec<Arc<dyn AgentTool>>, saved_prompt: String) {
160        self.state.tools = saved_tools;
161        self.state.system_prompt = saved_prompt;
162    }
163}