1081 lines
36 KiB
Rust
1081 lines
36 KiB
Rust
use std::sync::Arc;
|
||
use std::time::Duration;
|
||
|
||
use base64::Engine;
|
||
use lazyboy_contracts::ModelProvider;
|
||
use lazyboy_harness::{
|
||
CredentialChain, DynModel, ResolveModelRequest, connect_model, resolve_backend,
|
||
};
|
||
use rig_core::completion::message::{
|
||
AssistantContent, ImageDetail, ImageMediaType, Message, ToolResultContent, UserContent,
|
||
};
|
||
use rig_core::completion::{CompletionModel, ToolDefinition};
|
||
use serde_json::{Value, json};
|
||
use uuid::Uuid;
|
||
|
||
use crate::computer::{self, adapter_context_for};
|
||
use crate::db::{Actor, parse_mode};
|
||
use crate::state::AppState;
|
||
use crate::tools::{ToolCtx, dispatch, tool_definitions};
|
||
|
||
const SCREENSHOT_CAPTION: &str = "Desktop screenshot (1280x800) with yellow numbered marks. Click by those element ids. The live VNC view has no marks.";
|
||
|
||
const SYSTEM: &str = "You operate this bot's Linux desktop. The human always sees the live screen. You do not need a screenshot for every step.
|
||
|
||
Prefer the fast path, in this order:
|
||
1) shell, list_files, read_file, write_file
|
||
2) MCP tools when they match the task
|
||
3) browser for anything in Chromium: snapshot (page text + numbered elements), click/type/press by element id, navigate by URL. Do not pixel-click the Chromium window.
|
||
4) launch_app / open_path to open a site or file
|
||
5) computer_act only for native GUI that has no DOM (dialogs, canvas, XFCE)
|
||
|
||
When you use the browser tool:
|
||
- snapshot first; click {\"action\":\"click\",\"element\":N}; type {\"action\":\"type\",\"element\":N,\"text\":\"...\"}; open a URL with navigate.
|
||
- Yellow numbered marks on the screenshot match the element list. Click the number, not guessed pixels.
|
||
- If the control is not in the element list, login/2FA/CAPTCHA, or clicks do nothing, call request_takeover and stop. Do not guess-click.
|
||
|
||
computer_act examples (native windows only):
|
||
- {\"kind\":\"click\",\"element\":1}
|
||
- {\"kind\":\"click\",\"x\":N,\"y\":N}
|
||
- {\"kind\":\"type\",\"text\":\"...\"}
|
||
- {\"kind\":\"key\",\"key\":\"Return\"}
|
||
- {\"kind\":\"focus\",\"title\":\"Open File\"}
|
||
|
||
On a Team Computer, relative files live in your bot folder; use shared/ for shared work. Finish the user's task.";
|
||
|
||
pub async fn send(
|
||
state: &AppState,
|
||
actor: &Actor,
|
||
bot_id: &str,
|
||
thread_id: &str,
|
||
text: &str,
|
||
client_nonce: Option<&str>,
|
||
blocks: &[Value],
|
||
) -> Result<Value, String> {
|
||
let mut tx = state
|
||
.pool()
|
||
.begin()
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
let scoped: Option<(String, Option<String>)> = sqlx::query_as(
|
||
"SELECT t.title, t.room_id FROM threads t JOIN bots b ON b.id=t.bot_id
|
||
WHERE t.id=$1 AND t.bot_id=$2 AND t.space_id=$3 AND t.user_id=$4
|
||
AND b.space_id=$3 AND b.user_id=$4 AND t.status='active'
|
||
FOR UPDATE OF t",
|
||
)
|
||
.bind(thread_id)
|
||
.bind(bot_id)
|
||
.bind(&actor.space_id)
|
||
.bind(&actor.user_id)
|
||
.fetch_optional(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
let Some((current_title, room_id)) = scoped else {
|
||
return Err("session not found".into());
|
||
};
|
||
if let Some(nonce) = client_nonce {
|
||
let existing: Option<(String, Option<String>)> = sqlx::query_as(
|
||
"SELECT id, run_id FROM messages WHERE thread_id=$1 AND client_nonce=$2",
|
||
)
|
||
.bind(thread_id)
|
||
.bind(nonce)
|
||
.fetch_optional(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
if let Some((message_id, run_id)) = existing {
|
||
tx.rollback().await.map_err(|error| error.to_string())?;
|
||
return Ok(json!({
|
||
"messageId": message_id,
|
||
"runId": run_id,
|
||
"duplicate": true,
|
||
"queued": true
|
||
}));
|
||
}
|
||
}
|
||
let run_id = Uuid::new_v4().to_string();
|
||
let message_id = Uuid::new_v4().to_string();
|
||
let seq: i32 = sqlx::query_scalar(
|
||
"UPDATE threads
|
||
SET next_message_seq=next_message_seq+1, updated_at=now()
|
||
WHERE id=$1 RETURNING next_message_seq-1",
|
||
)
|
||
.bind(thread_id)
|
||
.fetch_one(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
sqlx::query(
|
||
"INSERT INTO messages (id,thread_id,seq,role,body,blocks,run_id,client_nonce)
|
||
VALUES ($1,$2,$3,'user',$4,$5,$6,$7)",
|
||
)
|
||
.bind(&message_id)
|
||
.bind(thread_id)
|
||
.bind(seq)
|
||
.bind(text)
|
||
.bind(json!(blocks))
|
||
.bind(&run_id)
|
||
.bind(client_nonce)
|
||
.execute(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
if crate::sessions::is_default_session_title(¤t_title) {
|
||
sqlx::query("UPDATE threads SET title=$2 WHERE id=$1")
|
||
.bind(thread_id)
|
||
.bind(crate::sessions::title_from_first_message(text))
|
||
.execute(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
}
|
||
let mut member_ids: Vec<String> = if let Some(room_id) = room_id.as_deref() {
|
||
sqlx::query_scalar("SELECT bot_id FROM room_members WHERE room_id=$1 ORDER BY created_at")
|
||
.bind(room_id)
|
||
.fetch_all(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?
|
||
} else {
|
||
vec![bot_id.to_string()]
|
||
};
|
||
if member_ids.is_empty() {
|
||
member_ids.push(bot_id.to_string());
|
||
}
|
||
for (index, member_id) in member_ids.iter().enumerate() {
|
||
let member_run = if index == 0 {
|
||
run_id.clone()
|
||
} else {
|
||
Uuid::new_v4().to_string()
|
||
};
|
||
sqlx::query(
|
||
"INSERT INTO runs (id,space_id,bot_id,thread_id,user_id,status,prompt,checkpoint)
|
||
VALUES ($1,$2,$3,$4,$5,'queued',$6,$7)",
|
||
)
|
||
.bind(&member_run)
|
||
.bind(&actor.space_id)
|
||
.bind(member_id)
|
||
.bind(&thread_id)
|
||
.bind(&actor.user_id)
|
||
.bind(text)
|
||
.bind(json!({"messageSeq":seq}))
|
||
.execute(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
}
|
||
let event_seq: i32 = sqlx::query_scalar(
|
||
"UPDATE threads SET next_event_seq=next_event_seq+1 WHERE id=$1 RETURNING next_event_seq",
|
||
)
|
||
.bind(thread_id)
|
||
.fetch_one(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
sqlx::query(
|
||
"INSERT INTO events (id,thread_id,seq,type,payload) VALUES ($1,$2,$3,'message.created',$4)",
|
||
)
|
||
.bind(Uuid::new_v4().to_string())
|
||
.bind(thread_id)
|
||
.bind(event_seq)
|
||
.bind(json!({"id":message_id,"seq":seq,"role":"user","body":text,"runId":run_id}))
|
||
.execute(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
let queued_behind_active: bool = sqlx::query_scalar(
|
||
"SELECT EXISTS(SELECT 1 FROM runs WHERE bot_id=$1 AND id<>$2
|
||
AND status IN ('queued','leased','running','waiting_input','waiting_takeover'))",
|
||
)
|
||
.bind(bot_id)
|
||
.bind(&run_id)
|
||
.fetch_one(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
if let Some(room_id) = room_id.as_deref() {
|
||
sqlx::query("UPDATE rooms SET updated_at=now() WHERE id=$1")
|
||
.bind(room_id)
|
||
.execute(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
}
|
||
tx.commit().await.map_err(|error| error.to_string())?;
|
||
Ok(json!({
|
||
"messageId": message_id,
|
||
"runId": run_id,
|
||
"duplicate": false,
|
||
"queued": true,
|
||
"queuedBehindActive": queued_behind_active
|
||
}))
|
||
}
|
||
|
||
pub async fn worker_loop(state: AppState) {
|
||
let inflight = Arc::new(tokio::sync::Semaphore::new(16));
|
||
let lease_owner = format!("api-{}", Uuid::new_v4());
|
||
loop {
|
||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||
let Ok(permit) = inflight.clone().try_acquire_owned() else {
|
||
continue;
|
||
};
|
||
let queued: Result<Option<(String, String, String, String, String, String)>, _> =
|
||
sqlx::query_as(
|
||
"WITH candidate AS (
|
||
SELECT r.id
|
||
FROM runs r
|
||
WHERE r.retry_count < r.max_retries
|
||
AND (
|
||
r.status='queued'
|
||
OR (
|
||
r.status IN ('leased','running')
|
||
AND (r.lease_expires_at IS NULL OR r.lease_expires_at < now())
|
||
)
|
||
)
|
||
AND NOT EXISTS (
|
||
SELECT 1 FROM runs a
|
||
WHERE a.bot_id=r.bot_id AND a.id<>r.id
|
||
AND a.status IN ('leased','running','waiting_input','waiting_takeover')
|
||
AND (a.lease_expires_at IS NULL OR a.lease_expires_at >= now())
|
||
)
|
||
ORDER BY CASE WHEN r.status='queued' THEN 1 ELSE 0 END, r.created_at
|
||
FOR UPDATE SKIP LOCKED
|
||
LIMIT 1
|
||
)
|
||
UPDATE runs r
|
||
SET status='leased', lease_owner=$1,
|
||
lease_expires_at=now()+interval '5 minutes',
|
||
lease_fence=lease_fence+1, retry_count=retry_count+1, updated_at=now()
|
||
FROM candidate c WHERE r.id=c.id
|
||
RETURNING r.id,r.bot_id,r.thread_id,r.prompt,r.user_id,r.space_id",
|
||
)
|
||
.bind(&lease_owner)
|
||
.fetch_optional(state.pool())
|
||
.await;
|
||
let Ok(Some((run_id, bot_id, thread_id, prompt, user_id, space_id))) = queued else {
|
||
drop(permit);
|
||
continue;
|
||
};
|
||
let state = state.clone();
|
||
let owner = lease_owner.clone();
|
||
tokio::spawn(async move {
|
||
let _permit = permit;
|
||
let actor = Actor { user_id, space_id };
|
||
if let Err(error) = execute_run(
|
||
&state, &actor, &owner, &run_id, &bot_id, &thread_id, &prompt,
|
||
)
|
||
.await
|
||
{
|
||
tracing::error!("run {run_id} failed: {error}");
|
||
let next_status: Option<String> = sqlx::query_scalar(
|
||
"UPDATE runs
|
||
SET status=CASE WHEN retry_count < max_retries THEN 'queued' ELSE 'failed' END,
|
||
error=$2, completed_at=CASE WHEN retry_count < max_retries THEN NULL ELSE now() END,
|
||
lease_owner=NULL, lease_expires_at=NULL, updated_at=now()
|
||
WHERE id=$1 AND lease_owner=$3 AND status IN ('leased','running') RETURNING status",
|
||
)
|
||
.bind(&run_id)
|
||
.bind(&error)
|
||
.bind(&owner)
|
||
.fetch_optional(state.pool())
|
||
.await
|
||
.ok()
|
||
.flatten();
|
||
if next_status.as_deref() == Some("failed") {
|
||
let _ = append_bot_message(
|
||
&state,
|
||
&thread_id,
|
||
&run_id,
|
||
&bot_id,
|
||
&format!("Run failed after retries: {error}"),
|
||
)
|
||
.await;
|
||
let _ = crate::sessions::append_event(
|
||
&state,
|
||
&thread_id,
|
||
"run.failed",
|
||
json!({"runId":run_id,"error":error}),
|
||
)
|
||
.await;
|
||
}
|
||
let _ = computer::release_screen_execution(&state, &run_id).await;
|
||
let _ = sqlx::query(
|
||
"UPDATE computers SET execution_bot_id = NULL, execution_run_id = NULL, execution_lease_expires_at = NULL, updated_at = now()
|
||
WHERE execution_run_id = $1",
|
||
)
|
||
.bind(&run_id)
|
||
.execute(state.pool())
|
||
.await;
|
||
}
|
||
});
|
||
}
|
||
}
|
||
|
||
async fn execute_run(
|
||
state: &AppState,
|
||
actor: &Actor,
|
||
lease_owner: &str,
|
||
run_id: &str,
|
||
bot_id: &str,
|
||
thread_id: &str,
|
||
prompt: &str,
|
||
) -> Result<(), String> {
|
||
let started = sqlx::query(
|
||
"UPDATE runs SET status='running', started_at=COALESCE(started_at,now()), updated_at=now()
|
||
WHERE id=$1 AND lease_owner=$2 AND lease_expires_at>now() AND status='leased'",
|
||
)
|
||
.bind(run_id)
|
||
.bind(lease_owner)
|
||
.execute(state.pool())
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
if started.rows_affected() != 1 {
|
||
let status: Option<String> = sqlx::query_scalar("SELECT status FROM runs WHERE id=$1")
|
||
.bind(run_id)
|
||
.fetch_optional(state.pool())
|
||
.await
|
||
.ok()
|
||
.flatten();
|
||
if halt_from_status(status.as_deref()).is_some() {
|
||
return Ok(());
|
||
}
|
||
return Err("run lease was lost before execution".into());
|
||
}
|
||
let _ = crate::sessions::append_event(state, thread_id, "run.started", json!({"runId":run_id}))
|
||
.await;
|
||
|
||
computer::boot(state, actor, bot_id).await?;
|
||
let bot = state
|
||
.db
|
||
.get_bot(actor, bot_id)
|
||
.await
|
||
.map_err(|error| error.to_string())?
|
||
.ok_or_else(|| "bot not found".to_string())?;
|
||
let computer = state
|
||
.db
|
||
.get_computer(bot.computer_id.as_deref().unwrap_or(""))
|
||
.await
|
||
.map_err(|error| error.to_string())?
|
||
.ok_or_else(|| "computer not found".to_string())?;
|
||
let computer_ref =
|
||
computer::computer_ref(&computer).ok_or_else(|| "computer is not running".to_string())?;
|
||
let bound = computer::ensure_bot_screen(state, actor, bot_id, &computer, Some(run_id)).await?;
|
||
let mut gui_block = bound.gui_block;
|
||
let screen = if let Some(row) = bound.row {
|
||
let row = computer::take_screen_execution(state, &row, run_id).await?;
|
||
if gui_block.is_none() {
|
||
gui_block =
|
||
computer::take_profile_lock(state, &computer, bot_id, &bot.name, run_id, &row)
|
||
.await?;
|
||
}
|
||
Some(row)
|
||
} else {
|
||
None
|
||
};
|
||
|
||
let space = state
|
||
.db
|
||
.get_space(actor)
|
||
.await
|
||
.map_err(|error| error.to_string())?
|
||
.ok_or_else(|| "workspace not found".to_string())?;
|
||
let provider = bot
|
||
.model_provider
|
||
.as_deref()
|
||
.or(Some(space.default_model_provider.as_str()))
|
||
.unwrap_or("xai")
|
||
.parse::<ModelProvider>()
|
||
.map_err(|error| error.to_string())?;
|
||
let model_id = bot
|
||
.model_id
|
||
.clone()
|
||
.filter(|value| !value.is_empty())
|
||
.or_else(|| Some(space.default_model_id.clone()).filter(|value| !value.is_empty()));
|
||
let backend = resolve_backend(ResolveModelRequest {
|
||
provider,
|
||
model_id,
|
||
base_url: space.default_model_base_url.clone(),
|
||
credentials: CredentialChain {
|
||
bot: None,
|
||
space: space.default_model_api_key.clone(),
|
||
env: lazyboy_harness::credential_from_env(provider),
|
||
},
|
||
})
|
||
.map_err(|error| error.to_string())?;
|
||
let model = connect_model(&backend).map_err(|error| error.to_string())?;
|
||
|
||
let ctx = Arc::new(ToolCtx {
|
||
sandbox: state.sandbox.clone(),
|
||
computer: computer_ref,
|
||
context: adapter_context_for(actor, bot_id, "run", screen.as_ref(), Some(run_id)),
|
||
mode: parse_mode(&computer.scope),
|
||
bot_id: bot_id.to_string(),
|
||
vision: backend.capabilities.vision,
|
||
gui_block,
|
||
previous_frame: std::sync::Mutex::new(None),
|
||
elements: std::sync::Mutex::new(Vec::new()),
|
||
miss_streak: std::sync::Mutex::new(0),
|
||
click_misses: std::sync::Mutex::new(0),
|
||
takeover_requested: std::sync::Mutex::new(false),
|
||
pool: state.pool().clone(),
|
||
memory: state.memory.clone(),
|
||
actor: actor.clone(),
|
||
session_id: thread_id.to_string(),
|
||
run_id: run_id.to_string(),
|
||
memory_enabled: bot.memory_enabled && state.memory.globally_enabled(),
|
||
mcp: state.mcp.clone(),
|
||
});
|
||
|
||
let mut defs = tool_definitions(ctx.memory_enabled);
|
||
let mcp_defs = state.mcp.definitions().await;
|
||
if !mcp_defs.is_empty() {
|
||
defs.extend(mcp_defs);
|
||
}
|
||
let (summary, summary_seq): (String, i32) = sqlx::query_as(
|
||
"SELECT history_summary, history_summary_seq FROM threads
|
||
WHERE id=$1 AND space_id=$2 AND user_id=$3",
|
||
)
|
||
.bind(thread_id)
|
||
.bind(&actor.space_id)
|
||
.bind(&actor.user_id)
|
||
.fetch_one(state.pool())
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
let checkpoint: Value = sqlx::query_scalar("SELECT checkpoint FROM runs WHERE id=$1")
|
||
.bind(run_id)
|
||
.fetch_one(state.pool())
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
let current_seq: i32 = checkpoint
|
||
.get("messageSeq")
|
||
.and_then(Value::as_i64)
|
||
.map(|seq| seq as i32)
|
||
.unwrap_or(i32::MAX);
|
||
let resume_after_takeover = checkpoint
|
||
.get("resumeAfterTakeover")
|
||
.and_then(Value::as_bool)
|
||
.unwrap_or(false);
|
||
if resume_after_takeover {
|
||
let _ = sqlx::query(
|
||
"UPDATE runs SET checkpoint = checkpoint - 'resumeAfterTakeover' WHERE id=$1",
|
||
)
|
||
.bind(run_id)
|
||
.execute(state.pool())
|
||
.await;
|
||
}
|
||
let history_end = if resume_after_takeover {
|
||
i32::MAX
|
||
} else {
|
||
current_seq
|
||
};
|
||
let recent: Vec<(String, String, Option<String>, Option<String>)> = sqlx::query_as(
|
||
"SELECT role, body, speaker_bot_id, speaker_name FROM (
|
||
SELECT m.role, m.body, m.seq, m.speaker_bot_id, b.name AS speaker_name
|
||
FROM messages m
|
||
LEFT JOIN bots b ON b.id=m.speaker_bot_id
|
||
WHERE m.thread_id=$1 AND m.seq>$2 AND m.seq<$3
|
||
ORDER BY m.seq DESC LIMIT 30
|
||
) history ORDER BY seq ASC",
|
||
)
|
||
.bind(thread_id)
|
||
.bind(history_window_start(summary_seq, current_seq))
|
||
.bind(history_end)
|
||
.fetch_all(state.pool())
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
let mut history: Vec<Message> = Vec::new();
|
||
if !summary.trim().is_empty() {
|
||
history.push(Message::User {
|
||
content: vec![UserContent::text(format!(
|
||
"Conversation summary through message {summary_seq}:\n{summary}"
|
||
))],
|
||
});
|
||
}
|
||
for (role, body, speaker_id, speaker_name) in recent {
|
||
if role == "user" {
|
||
history.push(Message::User {
|
||
content: vec![UserContent::text(body)],
|
||
});
|
||
} else if speaker_id.as_deref() == Some(bot_id) || speaker_id.is_none() {
|
||
history.push(Message::Assistant {
|
||
id: None,
|
||
content: vec![AssistantContent::text(body)],
|
||
});
|
||
} else {
|
||
let name = speaker_name.unwrap_or_else(|| "agent".into());
|
||
history.push(Message::User {
|
||
content: vec![UserContent::text(format!("[{name}]: {body}"))],
|
||
});
|
||
}
|
||
}
|
||
let mut first = if resume_after_takeover {
|
||
vec![UserContent::text(
|
||
"The user finished collaborating and released control. Continue the original task from the CURRENT screen. Do not restart from scratch.",
|
||
)]
|
||
} else {
|
||
vec![UserContent::text(prompt)]
|
||
};
|
||
let mut screenshots: u32 = 0;
|
||
let mut screenshot_bytes: u64 = 0;
|
||
if resume_after_takeover && ctx.gui_block.is_none() {
|
||
let outcome = dispatch(&ctx, "computer_observe", &json!({})).await;
|
||
first.push(UserContent::text(outcome.text));
|
||
if let Some(image) = outcome.image {
|
||
screenshot_bytes += image.len() as u64;
|
||
screenshots += 1;
|
||
first.extend(screenshot_parts(image));
|
||
}
|
||
}
|
||
let mut pending = Message::User { content: first };
|
||
let mut final_text = String::new();
|
||
let mut turns: u32 = 0;
|
||
let mut used_gui = false;
|
||
let memory = if ctx.memory_enabled {
|
||
match state
|
||
.memory
|
||
.recall(state.pool(), actor, bot_id, prompt, None)
|
||
.await
|
||
{
|
||
Ok(items) => state.memory.durable_block(&items),
|
||
Err(error) => {
|
||
tracing::warn!("memory retrieval failed for run {run_id}: {error}");
|
||
String::new()
|
||
}
|
||
}
|
||
} else {
|
||
String::new()
|
||
};
|
||
let mut preamble = if bot.instructions.trim().is_empty() {
|
||
SYSTEM.to_string()
|
||
} else {
|
||
format!(
|
||
"{SYSTEM}\n\nBot-specific instructions:\n{}",
|
||
bot.instructions.trim()
|
||
)
|
||
};
|
||
let room_mates: Vec<String> = sqlx::query_scalar(
|
||
"SELECT b.name FROM threads t
|
||
JOIN room_members m ON m.room_id=t.room_id
|
||
JOIN bots b ON b.id=m.bot_id
|
||
WHERE t.id=$1 AND b.id<>$2
|
||
ORDER BY b.name",
|
||
)
|
||
.bind(thread_id)
|
||
.bind(bot_id)
|
||
.fetch_all(state.pool())
|
||
.await
|
||
.unwrap_or_default();
|
||
if !room_mates.is_empty() {
|
||
preamble.push_str(&format!(
|
||
"\n\nYou are {} in a group chat with: {}. Reply as yourself only. Other agents' lines are prefixed with [Name]. Do not speak for them.",
|
||
bot.name,
|
||
room_mates.join("、")
|
||
));
|
||
}
|
||
let mcp_names: Vec<String> = defs
|
||
.iter()
|
||
.filter(|tool| tool.name.starts_with("mcp_"))
|
||
.map(|tool| tool.name.clone())
|
||
.collect();
|
||
if !mcp_names.is_empty() {
|
||
preamble.push_str(&format!(
|
||
"\n\nMCP tools available: {}. Use them when they help complete the user's request.",
|
||
mcp_names.join(", ")
|
||
));
|
||
}
|
||
if !memory.is_empty() {
|
||
preamble.push_str("\n\n");
|
||
preamble.push_str(&memory);
|
||
}
|
||
|
||
for _ in 0..24 {
|
||
turns += 1;
|
||
if let Some(halt) = renew_or_halt(state, run_id, lease_owner).await? {
|
||
return finish_halt(
|
||
state,
|
||
thread_id,
|
||
run_id,
|
||
halt,
|
||
turns,
|
||
screenshots,
|
||
screenshot_bytes,
|
||
&ctx,
|
||
used_gui,
|
||
)
|
||
.await;
|
||
}
|
||
drop_history_screenshots(&mut history);
|
||
let content = tokio::select! {
|
||
halt = wait_for_halt(state, run_id) => {
|
||
return finish_halt(
|
||
state,
|
||
thread_id,
|
||
run_id,
|
||
halt,
|
||
turns,
|
||
screenshots,
|
||
screenshot_bytes,
|
||
&ctx,
|
||
used_gui,
|
||
)
|
||
.await;
|
||
}
|
||
result = complete_once(&model, pending.clone(), &preamble, &history, &defs) => result
|
||
}?;
|
||
let assistant = Message::Assistant {
|
||
id: None,
|
||
content: content.clone(),
|
||
};
|
||
history.push(pending.clone());
|
||
history.push(assistant);
|
||
|
||
let mut calls = Vec::new();
|
||
for item in &content {
|
||
match item {
|
||
AssistantContent::Text(text) => final_text.push_str(&text.text),
|
||
AssistantContent::ToolCall(call) => calls.push(call.clone()),
|
||
_ => {}
|
||
}
|
||
}
|
||
if calls.is_empty() {
|
||
break;
|
||
}
|
||
final_text.clear();
|
||
let mut results = Vec::new();
|
||
let mut screen: Option<Vec<u8>> = None;
|
||
for call in calls {
|
||
if let Some(halt) = renew_or_halt(state, run_id, lease_owner).await? {
|
||
return finish_halt(
|
||
state,
|
||
thread_id,
|
||
run_id,
|
||
halt,
|
||
turns,
|
||
screenshots,
|
||
screenshot_bytes,
|
||
&ctx,
|
||
used_gui,
|
||
)
|
||
.await;
|
||
}
|
||
let name = call.function.name.clone();
|
||
used_gui |= matches!(
|
||
name.as_str(),
|
||
"computer_observe" | "computer_act" | "open_path" | "launch_app" | "browser"
|
||
);
|
||
let outcome = tokio::select! {
|
||
halt = wait_for_halt(state, run_id) => {
|
||
return finish_halt(
|
||
state,
|
||
thread_id,
|
||
run_id,
|
||
halt,
|
||
turns,
|
||
screenshots,
|
||
screenshot_bytes,
|
||
&ctx,
|
||
used_gui,
|
||
)
|
||
.await;
|
||
}
|
||
outcome = tokio::time::timeout(
|
||
Duration::from_secs(90),
|
||
dispatch(&ctx, &name, &call.function.arguments),
|
||
) => match outcome {
|
||
Ok(outcome) => outcome,
|
||
Err(_) => crate::tools::ToolOutcome {
|
||
text: format!("工具 {name} 執行逾時（90 秒），請稍後重試。"),
|
||
image: None,
|
||
pause: false,
|
||
},
|
||
}
|
||
};
|
||
// xAI rejects images inside tool results. Attach a changed
|
||
// screenshot as a following user image instead.
|
||
if let Some(image) = outcome.image {
|
||
screen = Some(image);
|
||
}
|
||
results.push(UserContent::tool_result_for(
|
||
call.id.clone(),
|
||
call.provider.clone(),
|
||
name,
|
||
vec![ToolResultContent::text(&outcome.text)],
|
||
));
|
||
if outcome.pause {
|
||
sqlx::query(
|
||
"UPDATE runs SET status = 'waiting_takeover',
|
||
checkpoint = COALESCE(checkpoint, '{}'::jsonb)
|
||
|| jsonb_build_object('resumeAfterTakeover', true),
|
||
updated_at = now()
|
||
WHERE id = $1",
|
||
)
|
||
.bind(run_id)
|
||
.execute(state.pool())
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
append_bot_message(state, thread_id, run_id, bot_id, &outcome.text).await?;
|
||
let click_misses = *ctx.click_misses.lock().unwrap();
|
||
record_run_metrics(
|
||
state,
|
||
thread_id,
|
||
run_id,
|
||
"run.paused",
|
||
turns,
|
||
screenshots,
|
||
screenshot_bytes,
|
||
click_misses,
|
||
used_gui,
|
||
true,
|
||
)
|
||
.await;
|
||
return Ok(());
|
||
}
|
||
}
|
||
if let Some(png) = screen {
|
||
screenshot_bytes += png.len() as u64;
|
||
screenshots += 1;
|
||
results.extend(screenshot_parts(png));
|
||
}
|
||
pending = Message::User { content: results };
|
||
}
|
||
|
||
let status: Option<String> = sqlx::query_scalar("SELECT status FROM runs WHERE id = $1")
|
||
.bind(run_id)
|
||
.fetch_optional(state.pool())
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
if let Some(halt) = halt_from_status(status.as_deref()) {
|
||
return finish_halt(
|
||
state,
|
||
thread_id,
|
||
run_id,
|
||
halt,
|
||
turns,
|
||
screenshots,
|
||
screenshot_bytes,
|
||
&ctx,
|
||
used_gui,
|
||
)
|
||
.await;
|
||
}
|
||
append_bot_message(state, thread_id, run_id, bot_id, &final_text).await?;
|
||
let completed = sqlx::query(
|
||
"UPDATE runs
|
||
SET status='completed', completed_at=now(), updated_at=now(),
|
||
lease_owner=NULL, lease_expires_at=NULL
|
||
WHERE id=$1 AND lease_owner=$2 AND status='running'",
|
||
)
|
||
.bind(run_id)
|
||
.bind(lease_owner)
|
||
.execute(state.pool())
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
if completed.rows_affected() != 1 {
|
||
return Err("run lease was lost before completion".into());
|
||
}
|
||
let click_misses = *ctx.click_misses.lock().unwrap();
|
||
let takeover = *ctx.takeover_requested.lock().unwrap();
|
||
record_run_metrics(
|
||
state,
|
||
thread_id,
|
||
run_id,
|
||
"run.completed",
|
||
turns,
|
||
screenshots,
|
||
screenshot_bytes,
|
||
click_misses,
|
||
used_gui,
|
||
takeover,
|
||
)
|
||
.await;
|
||
computer::release_screen_execution(state, run_id).await?;
|
||
sqlx::query(
|
||
"UPDATE computers SET execution_bot_id = NULL, execution_run_id = NULL, execution_lease_expires_at = NULL, updated_at = now()
|
||
WHERE execution_run_id = $1",
|
||
)
|
||
.bind(run_id)
|
||
.execute(state.pool())
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
Ok(())
|
||
}
|
||
|
||
async fn complete_once(
|
||
model: &DynModel,
|
||
pending: Message,
|
||
preamble: &str,
|
||
history: &[Message],
|
||
defs: &[ToolDefinition],
|
||
) -> Result<Vec<AssistantContent>, String> {
|
||
match model {
|
||
DynModel::Xai(model) => complete_with(model, pending, preamble, history, defs).await,
|
||
DynModel::OpenAi(model) => complete_with(model, pending, preamble, history, defs).await,
|
||
}
|
||
}
|
||
|
||
async fn complete_with<M>(
|
||
model: &M,
|
||
pending: Message,
|
||
preamble: &str,
|
||
history: &[Message],
|
||
defs: &[ToolDefinition],
|
||
) -> Result<Vec<AssistantContent>, String>
|
||
where
|
||
M: CompletionModel + Clone,
|
||
{
|
||
let request = model
|
||
.completion_request(pending)
|
||
.preamble(preamble.to_string())
|
||
.messages(history.to_vec())
|
||
.tools(defs.to_vec())
|
||
.build();
|
||
let response = tokio::time::timeout(Duration::from_secs(120), model.completion(request))
|
||
.await
|
||
.map_err(|_| "AI 回應逾時（120 秒）".to_string())?
|
||
.map_err(|error| error.to_string())?;
|
||
Ok(response.choice.into_iter().collect())
|
||
}
|
||
|
||
fn screenshot_parts(image: Vec<u8>) -> Vec<UserContent> {
|
||
let media = if image.starts_with(&[0xFF, 0xD8, 0xFF]) {
|
||
ImageMediaType::JPEG
|
||
} else {
|
||
ImageMediaType::PNG
|
||
};
|
||
let encoded = base64::engine::general_purpose::STANDARD.encode(image);
|
||
vec![
|
||
UserContent::text(SCREENSHOT_CAPTION),
|
||
UserContent::image_base64(encoded, Some(media), Some(ImageDetail::Low)),
|
||
]
|
||
}
|
||
|
||
fn drop_history_screenshots(history: &mut [Message]) {
|
||
for message in history.iter_mut() {
|
||
let Message::User { content } = message else {
|
||
continue;
|
||
};
|
||
content.retain(|part| match part {
|
||
UserContent::Image(_) => false,
|
||
UserContent::Text(text) if text.text == SCREENSHOT_CAPTION => false,
|
||
_ => true,
|
||
});
|
||
}
|
||
}
|
||
|
||
pub(crate) async fn append_bot_message(
|
||
state: &AppState,
|
||
thread_id: &str,
|
||
run_id: &str,
|
||
bot_id: &str,
|
||
body: &str,
|
||
) -> Result<(), String> {
|
||
let mut tx = state
|
||
.pool()
|
||
.begin()
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
let seq: i32 = sqlx::query_scalar(
|
||
"UPDATE threads SET next_message_seq=next_message_seq+1,updated_at=now()
|
||
WHERE id=$1 RETURNING next_message_seq-1",
|
||
)
|
||
.bind(thread_id)
|
||
.fetch_one(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
let message_id = Uuid::new_v4().to_string();
|
||
sqlx::query(
|
||
"INSERT INTO messages (id,thread_id,seq,role,body,run_id,speaker_bot_id)
|
||
VALUES ($1,$2,$3,'assistant',$4,$5,$6)",
|
||
)
|
||
.bind(&message_id)
|
||
.bind(thread_id)
|
||
.bind(seq)
|
||
.bind(body)
|
||
.bind(run_id)
|
||
.bind(bot_id)
|
||
.execute(&mut *tx)
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
tx.commit().await.map_err(|error| error.to_string())?;
|
||
let _ = crate::sessions::append_event(
|
||
state,
|
||
thread_id,
|
||
"message.created",
|
||
json!({"id":message_id,"seq":seq,"role":"assistant","body":body,"runId":run_id}),
|
||
)
|
||
.await;
|
||
Ok(())
|
||
}
|
||
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
enum RunHalt {
|
||
Cancelled,
|
||
Takeover,
|
||
}
|
||
|
||
fn halt_from_status(status: Option<&str>) -> Option<RunHalt> {
|
||
match status {
|
||
Some("cancelled") => Some(RunHalt::Cancelled),
|
||
Some("waiting_takeover") => Some(RunHalt::Takeover),
|
||
_ => None,
|
||
}
|
||
}
|
||
|
||
async fn run_status_halt(state: &AppState, run_id: &str) -> Result<Option<RunHalt>, String> {
|
||
let status: Option<String> = sqlx::query_scalar("SELECT status FROM runs WHERE id = $1")
|
||
.bind(run_id)
|
||
.fetch_optional(state.pool())
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
Ok(halt_from_status(status.as_deref()))
|
||
}
|
||
|
||
async fn wait_for_halt(state: &AppState, run_id: &str) -> RunHalt {
|
||
loop {
|
||
if let Ok(Some(halt)) = run_status_halt(state, run_id).await {
|
||
return halt;
|
||
}
|
||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||
}
|
||
}
|
||
|
||
async fn renew_or_halt(
|
||
state: &AppState,
|
||
run_id: &str,
|
||
lease_owner: &str,
|
||
) -> Result<Option<RunHalt>, String> {
|
||
if let Some(halt) = run_status_halt(state, run_id).await? {
|
||
return Ok(Some(halt));
|
||
}
|
||
match renew_lease(state, run_id, lease_owner).await {
|
||
Ok(()) => run_status_halt(state, run_id).await,
|
||
Err(error) => {
|
||
if let Some(halt) = run_status_halt(state, run_id).await? {
|
||
Ok(Some(halt))
|
||
} else {
|
||
Err(error)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn finish_halt(
|
||
state: &AppState,
|
||
thread_id: &str,
|
||
run_id: &str,
|
||
halt: RunHalt,
|
||
turns: u32,
|
||
screenshots: u32,
|
||
screenshot_bytes: u64,
|
||
ctx: &crate::tools::ToolCtx,
|
||
used_gui: bool,
|
||
) -> Result<(), String> {
|
||
match halt {
|
||
RunHalt::Cancelled => Ok(()),
|
||
RunHalt::Takeover => {
|
||
let _ = sqlx::query(
|
||
"UPDATE runs SET lease_owner=NULL, lease_expires_at=NULL, updated_at=now()
|
||
WHERE id=$1 AND status='waiting_takeover'",
|
||
)
|
||
.bind(run_id)
|
||
.execute(state.pool())
|
||
.await;
|
||
let click_misses = *ctx.click_misses.lock().unwrap();
|
||
record_run_metrics(
|
||
state,
|
||
thread_id,
|
||
run_id,
|
||
"run.paused",
|
||
turns,
|
||
screenshots,
|
||
screenshot_bytes,
|
||
click_misses,
|
||
used_gui,
|
||
true,
|
||
)
|
||
.await;
|
||
Ok(())
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn renew_lease(state: &AppState, run_id: &str, lease_owner: &str) -> Result<(), String> {
|
||
let renewed = sqlx::query(
|
||
"UPDATE runs SET lease_expires_at=now()+interval '5 minutes',updated_at=now()
|
||
WHERE id=$1 AND lease_owner=$2 AND status IN ('leased','running')",
|
||
)
|
||
.bind(run_id)
|
||
.bind(lease_owner)
|
||
.execute(state.pool())
|
||
.await
|
||
.map_err(|error| error.to_string())?;
|
||
if renewed.rows_affected() == 1 {
|
||
Ok(())
|
||
} else {
|
||
Err("run lease was lost".into())
|
||
}
|
||
}
|
||
|
||
fn history_window_start(summary_seq: i32, current_seq: i32) -> i32 {
|
||
summary_seq.min(current_seq.saturating_sub(1)).max(0)
|
||
}
|
||
|
||
async fn record_run_metrics(
|
||
state: &AppState,
|
||
thread_id: &str,
|
||
run_id: &str,
|
||
event: &str,
|
||
turns: u32,
|
||
screenshots: u32,
|
||
screenshot_bytes: u64,
|
||
click_misses: u32,
|
||
used_gui: bool,
|
||
takeover: bool,
|
||
) {
|
||
let payload = json!({
|
||
"runId": run_id,
|
||
"turns": turns,
|
||
"screenshotsToModel": screenshots,
|
||
"screenshotBytes": screenshot_bytes,
|
||
"clickMisses": click_misses,
|
||
"usedGui": used_gui,
|
||
"takeover": takeover,
|
||
});
|
||
tracing::info!(
|
||
run_id,
|
||
turns,
|
||
screenshots,
|
||
screenshot_bytes,
|
||
click_misses,
|
||
used_gui,
|
||
takeover,
|
||
"run metrics"
|
||
);
|
||
let _ = crate::sessions::append_event(state, thread_id, event, payload).await;
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::{RunHalt, halt_from_status, history_window_start, screenshot_parts};
|
||
use rig_core::completion::message::UserContent;
|
||
|
||
#[test]
|
||
fn history_never_reads_past_the_current_prompt() {
|
||
assert_eq!(history_window_start(10, 5), 4);
|
||
assert_eq!(history_window_start(3, 20), 3);
|
||
assert_eq!(history_window_start(-1, 1), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn screenshot_parts_keeps_caption_and_image() {
|
||
let parts = screenshot_parts(vec![0xFF, 0xD8, 0xFF, 0x00]);
|
||
assert_eq!(parts.len(), 2);
|
||
assert!(matches!(parts[0], UserContent::Text(_)));
|
||
assert!(matches!(parts[1], UserContent::Image(_)));
|
||
}
|
||
|
||
#[test]
|
||
fn halt_maps_paused_and_cancelled_runs() {
|
||
assert_eq!(
|
||
halt_from_status(Some("cancelled")),
|
||
Some(RunHalt::Cancelled)
|
||
);
|
||
assert_eq!(
|
||
halt_from_status(Some("waiting_takeover")),
|
||
Some(RunHalt::Takeover)
|
||
);
|
||
assert_eq!(halt_from_status(Some("running")), None);
|
||
assert_eq!(halt_from_status(None), None);
|
||
}
|
||
}
|