1use std::collections::HashMap;
2use std::sync::Arc;
3use std::time::Duration;
4
5use futures::StreamExt;
6use serde_json::Value;
7
8use crate::artifacts::ArtifactStore;
9use crate::host::{CodeModeHost, ExecCtx, StepDecision};
10use crate::pool::{PoolConfig, RunnerDisposition, RunnerPool, RunnerSpawn};
11use crate::protocol::{CodeModeRunnerInput, CodeModeRunnerOutput};
12use crate::runner_io::{decode_runner_output, terminate_code_mode_runner, write_runner_input};
13use crate::types::{CodeModeCaller, CodeModeExecutionResponse, CodeModeSurface, ToolScope, UiLink};
14use crate::{normalize_user_code, CodeModeConfig, ToolError};
15
16use super::budget::RunBudget;
17use super::proxy::{build_proxy, load_entries};
18use super::tool_dispatch::{handle_tool_call, ToolCallContext};
19use super::{finish_response, CodeModeExecutionOutcome};
20
21pub(crate) struct SubprocessExecution<'a, H: CodeModeHost> {
22 pub(crate) host: Option<&'a H>,
23 pub(crate) runner_pool: Option<&'a RunnerPool>,
24 pub(crate) code: &'a str,
25 pub(crate) caller: CodeModeCaller,
26 pub(crate) surface: CodeModeSurface,
27 pub(crate) config: CodeModeConfig,
28 pub(crate) scope: ToolScope,
29 pub(crate) execution_id: Option<Arc<str>>,
30 pub(crate) ui_capture: Arc<std::sync::Mutex<Option<UiLink>>>,
31}
32
33pub(crate) async fn execute_in_subprocess<H: CodeModeHost>(
34 request: SubprocessExecution<'_, H>,
35) -> Result<CodeModeExecutionOutcome, ToolError> {
36 let entries = load_entries(
37 request.host,
38 &request.caller,
39 request.surface,
40 &request.scope,
41 )
42 .await?;
43 let config = request.config;
44 let mut budget = RunBudget::new(&config);
45 let proxy = build_proxy(&entries, config.semantic_search.blend_weight)?;
46 let fallback_pool;
47 let pool = if let Some(pool) = request.runner_pool {
48 pool
49 } else {
50 fallback_pool = RunnerPool::new(
51 PoolConfig {
52 size: 0,
53 recycle_after: 1,
54 max_overflow: 1,
55 },
56 RunnerSpawn::current_exe()?,
57 );
58 &fallback_pool
59 };
60 let mut lease = pool.checkout().await?;
61 let deadline = tokio::time::Instant::now() + Duration::from_millis(config.timeout_ms.max(1));
62 write_with_deadline(
63 &mut lease.handle_mut()?.stdin,
64 &CodeModeRunnerInput::Start {
65 code: normalize_user_code(request.code),
66 proxy,
67 },
68 deadline,
69 )
70 .await?;
71
72 let mut calls = Vec::new();
73 let mut step_ordinals: HashMap<u64, (u64, String)> = HashMap::new();
74 let mut next_step_ordinal = 0u64;
75 let artifact_run_id = request
76 .execution_id
77 .as_deref()
78 .map(ToOwned::to_owned)
79 .unwrap_or_else(|| ulid::Ulid::generate().to_string());
80 let artifact_store = ArtifactStore::new(artifact_run_id)?;
81 crate::artifacts::prune::prune_old_runs(&crate::soma_home().join("code-mode-artifacts"), 256)
82 .await
83 .map_err(|err| ToolError::internal_message(format!("prune artifacts: {err}")))?;
84 let mut tool_ctx = ToolCallContext {
85 host: request.host,
86 entries: &entries,
87 caller: &request.caller,
88 surface: request.surface,
89 scope: &request.scope,
90 execution_id: &request.execution_id,
91 ui_capture: &request.ui_capture,
92 calls: &mut calls,
93 };
94
95 loop {
96 let output = next_output(lease.handle_mut()?, deadline).await?;
97 match output {
98 CodeModeRunnerOutput::ToolCall { seq, id, params } => {
99 let result = handle_tool_call(&mut tool_ctx, &mut budget, seq, id, params).await;
100 settle(seq, result, &mut lease.handle_mut()?.stdin, deadline).await?;
101 }
102 CodeModeRunnerOutput::ArtifactWrite {
103 seq,
104 path,
105 content,
106 content_type,
107 } => {
108 let result = match budget.record_operation("artifact write") {
109 Ok(()) => artifact_store
110 .write_text(&path, &content, content_type.as_deref())
111 .await
112 .and_then(to_value),
113 Err(error) => Err(error),
114 };
115 settle(seq, result, &mut lease.handle_mut()?.stdin, deadline).await?;
116 }
117 CodeModeRunnerOutput::SnippetResolve { seq, name, input } => {
118 let result = match budget.record_operation("snippet resolve") {
119 Ok(()) => resolve_snippet(request.host, name, input).await,
120 Err(error) => Err(error),
121 };
122 match result {
123 Ok((code, input)) => {
124 write_with_deadline(
125 &mut lease.handle_mut()?.stdin,
126 &CodeModeRunnerInput::SnippetResolved { seq, code, input },
127 deadline,
128 )
129 .await?;
130 }
131 Err(error) => {
132 write_error(seq, error, &mut lease.handle_mut()?.stdin, deadline).await?
133 }
134 }
135 }
136 CodeModeRunnerOutput::StepBegin { seq, name } => {
137 if let Err(error) = budget.record_operation("step") {
138 write_error(seq, error, &mut lease.handle_mut()?.stdin, deadline).await?;
139 continue;
140 }
141 let ordinal = next_step_ordinal;
142 next_step_ordinal = next_step_ordinal.saturating_add(1);
143 step_ordinals.insert(seq, (ordinal, name.clone()));
144 let decision = decide_step(
145 request.host,
146 request.execution_id.clone(),
147 seq,
148 ordinal,
149 &name,
150 )
151 .await;
152 match decision {
153 StepDecision::Replay(value) => {
154 write_with_deadline(
155 &mut lease.handle_mut()?.stdin,
156 &CodeModeRunnerInput::StepDecision {
157 seq,
158 replay: Some(value),
159 },
160 deadline,
161 )
162 .await?;
163 }
164 StepDecision::Execute => {
165 write_with_deadline(
166 &mut lease.handle_mut()?.stdin,
167 &CodeModeRunnerInput::StepDecision { seq, replay: None },
168 deadline,
169 )
170 .await?;
171 }
172 StepDecision::Error { kind, message } => {
173 write_with_deadline(
174 &mut lease.handle_mut()?.stdin,
175 &CodeModeRunnerInput::ToolError { seq, kind, message },
176 deadline,
177 )
178 .await?;
179 }
180 }
181 }
182 CodeModeRunnerOutput::StepResult { seq, value } => {
183 let result = record_step(
184 request.host,
185 request.execution_id.clone(),
186 seq,
187 &value,
188 &step_ordinals,
189 )
190 .await;
191 match result {
192 Ok(()) => {
193 write_with_deadline(
194 &mut lease.handle_mut()?.stdin,
195 &CodeModeRunnerInput::StepRecorded { seq },
196 deadline,
197 )
198 .await?;
199 }
200 Err(error) => {
201 write_error(seq, error, &mut lease.handle_mut()?.stdin, deadline).await?
202 }
203 }
204 }
205 CodeModeRunnerOutput::Done { result, logs } => {
206 lease.handle_mut()?.stderr.flush_settle().await;
207 let mut logs = logs;
208 logs.extend(lease.handle_mut()?.stderr.take_since_and_clear(0).await);
209 let logs = budget.cap_logs(logs);
210 let raw = CodeModeExecutionResponse {
211 result: result.into_response_result(),
212 calls,
213 logs,
214 error: None,
215 ui: request
216 .ui_capture
217 .lock()
218 .ok()
219 .and_then(|guard| guard.clone()),
220 };
221 let response = finish_response(raw, &config);
222 let handle = lease.handle_mut()?;
223 handle.success_count = handle.success_count.saturating_add(1);
224 let disposition = RunnerDisposition::from_success_count(
225 handle.success_count,
226 pool.config().recycle_after,
227 );
228 pool.release(lease, disposition).await;
229 return response;
230 }
231 CodeModeRunnerOutput::Error { kind, message } => {
232 return Err(ToolError::Sdk {
233 sdk_kind: kind,
234 message,
235 });
236 }
237 }
238 }
239}
240
241async fn resolve_snippet<H: CodeModeHost>(
242 host: Option<&H>,
243 name: String,
244 input: Value,
245) -> Result<(String, Value), ToolError> {
246 let host = host.ok_or_else(|| ToolError::UnknownInstance {
247 message: format!("unknown Code Mode snippet `{name}`"),
248 valid: Vec::new(),
249 })?;
250 let resolved = host.resolve_snippet(&name, input).await?;
251 Ok((resolved.code, resolved.input))
252}
253
254async fn decide_step<H: CodeModeHost>(
255 host: Option<&H>,
256 execution_id: Option<Arc<str>>,
257 seq: u64,
258 ordinal: u64,
259 name: &str,
260) -> StepDecision {
261 match host {
262 Some(host) => {
263 host.decide_step(
264 ExecCtx {
265 seq,
266 execution_id,
267 step_ordinal: Some(ordinal),
268 },
269 name,
270 )
271 .await
272 }
273 None => StepDecision::Execute,
274 }
275}
276
277async fn record_step<H: CodeModeHost>(
278 host: Option<&H>,
279 execution_id: Option<Arc<str>>,
280 seq: u64,
281 value: &Value,
282 step_ordinals: &HashMap<u64, (u64, String)>,
283) -> Result<(), ToolError> {
284 let Some(host) = host else {
285 return Ok(());
286 };
287 let (ordinal, name) = step_ordinals
288 .get(&seq)
289 .ok_or_else(|| ToolError::internal_message("runner returned an unknown step result seq"))?;
290 host.record_step(
291 ExecCtx {
292 seq,
293 execution_id,
294 step_ordinal: Some(*ordinal),
295 },
296 name,
297 value,
298 )
299 .await
300}
301
302async fn next_output(
303 runner: &mut crate::pool::RunnerHandle,
304 deadline: tokio::time::Instant,
305) -> Result<CodeModeRunnerOutput, ToolError> {
306 match tokio::time::timeout_at(deadline, runner.lines.next()).await {
307 Ok(Some(Ok(line))) => decode_runner_output(&line),
308 Ok(Some(Err(error))) => Err(ToolError::internal_message(format!(
309 "failed to read runner output: {error}"
310 ))),
311 Ok(None) => Err(ToolError::internal_message(
312 "runner exited before completion",
313 )),
314 Err(_) => {
315 terminate_code_mode_runner(&mut runner.child, runner.child_pid).await;
316 Err(ToolError::Sdk {
317 sdk_kind: "timeout".to_string(),
318 message: "Code Mode execution timed out".to_string(),
319 })
320 }
321 }
322}
323
324async fn settle<W: tokio::io::AsyncWriteExt + Unpin>(
325 seq: u64,
326 result: Result<Value, ToolError>,
327 writer: &mut W,
328 deadline: tokio::time::Instant,
329) -> Result<(), ToolError> {
330 match result {
331 Ok(result) => {
332 write_with_deadline(
333 writer,
334 &CodeModeRunnerInput::ToolResult { seq, result },
335 deadline,
336 )
337 .await
338 }
339 Err(error) => write_error(seq, error, writer, deadline).await,
340 }
341}
342
343async fn write_error<W: tokio::io::AsyncWriteExt + Unpin>(
344 seq: u64,
345 error: ToolError,
346 writer: &mut W,
347 deadline: tokio::time::Instant,
348) -> Result<(), ToolError> {
349 write_with_deadline(
350 writer,
351 &CodeModeRunnerInput::ToolError {
352 seq,
353 kind: error.kind().to_string(),
354 message: error.user_message().to_string(),
355 },
356 deadline,
357 )
358 .await
359}
360
361async fn write_with_deadline<W: tokio::io::AsyncWriteExt + Unpin>(
362 writer: &mut W,
363 input: &CodeModeRunnerInput,
364 deadline: tokio::time::Instant,
365) -> Result<(), ToolError> {
366 tokio::time::timeout_at(deadline, write_runner_input(writer, input))
367 .await
368 .map_err(|_| ToolError::Sdk {
369 sdk_kind: "timeout".to_string(),
370 message: "Code Mode runner write timed out".to_string(),
371 })?
372}
373
374fn to_value<T: serde::Serialize>(value: T) -> Result<Value, ToolError> {
375 serde_json::to_value(value).map_err(serialize_error)
376}
377
378fn serialize_error(error: serde_json::Error) -> ToolError {
379 ToolError::internal_message(format!("failed to serialize Code Mode value: {error}"))
380}