M2 audit: excise managed connectors and xAI media-gen tools
Managed connectors (grok.com MCP admin) removed root-and-branch: - The managed-MCP fetch/injection pipeline is gone, including the whole kigi-shell-session-support crate (managed-config fetch client, gateway tool catalog + dispatch, header injection, refresh task), reactive managed re-auth, mcp_doctor's grok.com-source discovery, and the [managed_mcps] config surface. - TUI: the 'Managed by grok.com' section, connectors URL/deep-link, Action::OpenManagedConnectors, and session_team_id are gone. Local MCP management (list/toggle/add/remove/auth/tools) is fully intact. - Kept as LOCAL policy: managed-settings.json MCP allow/deny enforcement, the multi-source local MCP merge, folder-trust gating. PluginOrigin Project/User labels kept (they tag locally discovered plugin dirs). imagine/media-gen tools (xAI image/video generation) removed: - image_gen, image_edit, video_gen, image_to_video, reference_to_video implementations, registrations, ToolKind/ToolInput/Output variants (serde-safe), config plumbing end to end, ZDR video machinery, /imagine + /imagine-video commands and guidance text, the bundled imagine skill (added to legacy cleanup so user installs delete it), and the media-gen render path. - Kept: image INPUT (paste/attach, [Image #N] meta, pdf/image fetch, clipboard wrap), generic media-ref rendering, and the generic tool 401-retry machinery (tests renamed, assertions unweakened). - deploy_app stays: it is a permanently-disabled local stub deploying nowhere. 121 files changed, 8 deleted. Gates: workspace check/clippy 0/0, fmt, deny ok; suites green (tools 2554, shell 4862, tui 6608, workspace 1042). Remaining grok.com strings live only in the auth-method ids and changelog archives (§9/M3 sweep).
This commit is contained in:
@@ -8,18 +8,12 @@ pub const SENT_BEARER_PREFIX_LEN: usize = 12;
|
||||
/// Which tool endpoint produced the 401.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ToolConsumer {
|
||||
ImageGen,
|
||||
VideoGenStart,
|
||||
VideoGenPoll,
|
||||
WebSearch,
|
||||
}
|
||||
|
||||
impl ToolConsumer {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::ImageGen => "ImageGen",
|
||||
Self::VideoGenStart => "VideoGen.start",
|
||||
Self::VideoGenPoll => "VideoGen.poll",
|
||||
Self::WebSearch => "WebSearch",
|
||||
}
|
||||
}
|
||||
@@ -95,9 +89,6 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn tool_consumer_as_str_stable_identifiers() {
|
||||
assert_eq!(ToolConsumer::ImageGen.as_str(), "ImageGen");
|
||||
assert_eq!(ToolConsumer::VideoGenStart.as_str(), "VideoGen.start");
|
||||
assert_eq!(ToolConsumer::VideoGenPoll.as_str(), "VideoGen.poll");
|
||||
assert_eq!(ToolConsumer::WebSearch.as_str(), "WebSearch");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,704 +0,0 @@
|
||||
//! `image_edit` tool — edits or transforms images via the xAI Imagine
|
||||
//! `/images/edits` endpoint using one or more reference images.
|
||||
//!
|
||||
//! Use cases include likeness preservation, style transfer, subject lock,
|
||||
//! remixing, and general image-to-image editing. The model chooses this
|
||||
//! tool (instead of `image_gen`) when the user provides reference photos.
|
||||
//!
|
||||
//! Reference images are specified as filesystem paths or
|
||||
//! `data:image/...;base64,...` URLs. The tool reads the bytes, compresses
|
||||
//! them to fit API limits, and POSTs to the edit endpoint.
|
||||
//!
|
||||
//! Shares the same [`ImageGenClient`] and session credentials as
|
||||
//! `image_gen` — no additional configuration is needed.
|
||||
|
||||
use std::io::Cursor;
|
||||
|
||||
use base64::Engine as _;
|
||||
use image::ImageReader;
|
||||
use reqwest::header::AUTHORIZATION;
|
||||
|
||||
use crate::attribution::ToolConsumer;
|
||||
use crate::implementations::grok_build::image_gen::{ImageGenClient, ImageGenResponse};
|
||||
use crate::types::output::{MediaGenOutput, ToolOutput};
|
||||
use crate::types::requirements::{Expr, ToolRequirement};
|
||||
use crate::types::resources::SessionFolder;
|
||||
use crate::types::tool::{ToolKind, ToolNamespace};
|
||||
use crate::util::image_compress::{FilterType, ReEncodeParams, re_encode_under_limit};
|
||||
|
||||
const XAI_IMAGINE_MODEL: &str = "grok-imagine-image-quality";
|
||||
|
||||
/// Size/dimension limits for reference images sent to the Imagine API.
|
||||
/// Tighter than the vision path; the backend returns 400 when exceeded.
|
||||
const MAX_REF_RAW_BYTES: usize = 400 * 1024;
|
||||
const MAX_REF_DIMENSION: u32 = 768;
|
||||
const MIN_REF_DIMENSION: u32 = 256;
|
||||
const REF_QUALITY_STEPS: &[u8] = &[80, 65, 50, 35];
|
||||
const MAX_REF_DECODE_PIXELS: u64 = 12_000_000;
|
||||
|
||||
pub const IMAGE_EDIT_TOOL_NAME: &str = "image_edit";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Compression
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Compress a reference image to fit within Imagine API limits.
|
||||
///
|
||||
/// Returns `(bytes, mime)`. Small JPEG/PNG inputs pass through unchanged.
|
||||
fn compress_reference(
|
||||
raw_bytes: Vec<u8>,
|
||||
) -> Result<(Vec<u8>, &'static str), kigi_tool_runtime::ToolError> {
|
||||
// Fast path: small JPEG/PNG passes through unchanged. Other formats
|
||||
// (WebP, GIF, etc.) always re-encode to guarantee API-compatible output.
|
||||
if raw_bytes.len() <= MAX_REF_RAW_BYTES
|
||||
&& let Some(kind) = infer::get(&raw_bytes)
|
||||
{
|
||||
match kind.mime_type() {
|
||||
"image/jpeg" => return Ok((raw_bytes, "image/jpeg")),
|
||||
"image/png" => return Ok((raw_bytes, "image/png")),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// Refuse to decode absurdly large images.
|
||||
let reader = ImageReader::new(Cursor::new(&raw_bytes))
|
||||
.with_guessed_format()
|
||||
.map_err(|_| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(
|
||||
"could not detect image format for reference",
|
||||
)
|
||||
})?;
|
||||
|
||||
if let Ok((w, h)) = reader.into_dimensions()
|
||||
&& (w as u64) * (h as u64) > MAX_REF_DECODE_PIXELS
|
||||
{
|
||||
return Err(kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"image reference is too large to process ({w}\u{00d7}{h} pixels)",
|
||||
)));
|
||||
}
|
||||
|
||||
// `into_dimensions` consumed the reader; re-open to decode.
|
||||
let img = ImageReader::new(Cursor::new(&raw_bytes))
|
||||
.with_guessed_format()
|
||||
.ok()
|
||||
.and_then(|r| r.decode().ok())
|
||||
.ok_or_else(|| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments("failed to decode image reference")
|
||||
})?;
|
||||
|
||||
let params = ReEncodeParams {
|
||||
max_bytes: MAX_REF_RAW_BYTES,
|
||||
max_side_px: MAX_REF_DIMENSION,
|
||||
// Imagine backend limits are side-based; no pixel-area cap applies.
|
||||
max_pixels: u64::MAX,
|
||||
min_side_px: MIN_REF_DIMENSION,
|
||||
quality_steps: REF_QUALITY_STEPS,
|
||||
filter: FilterType::Lanczos3,
|
||||
};
|
||||
|
||||
let (buf, _w, _h, mime) = re_encode_under_limit(&img, ¶ms).map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"could not compress image reference small enough for Imagine API: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok((buf, mime))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Reference resolution
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Resolve a reference (filesystem path or `data:image/...;base64,...` URL)
|
||||
/// into a compressed data URL for the Imagine API.
|
||||
async fn resolve_to_data_url(value: &str) -> Result<String, kigi_tool_runtime::ToolError> {
|
||||
let value = value.trim();
|
||||
// Accept `file://` URIs (e.g. an attachment's durable URI) by reading
|
||||
// the underlying path. Data URLs and bare paths are untouched.
|
||||
let value = value.strip_prefix("file://").unwrap_or(value);
|
||||
|
||||
let raw_bytes = if value.starts_with("data:image/") {
|
||||
let comma = value.find(',').ok_or_else(|| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments("malformed data URL in image reference")
|
||||
})?;
|
||||
if !value[..comma].contains(";base64") {
|
||||
return Err(kigi_tool_runtime::ToolError::invalid_arguments(
|
||||
"image references only support base64 data URLs",
|
||||
));
|
||||
}
|
||||
base64::engine::general_purpose::STANDARD
|
||||
.decode(&value[comma + 1..])
|
||||
.map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"invalid base64 in image reference: {e}"
|
||||
))
|
||||
})?
|
||||
} else {
|
||||
tokio::fs::read(value).await.map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"image reference not readable: {value} ({e})"
|
||||
))
|
||||
})?
|
||||
};
|
||||
|
||||
if raw_bytes.is_empty() {
|
||||
return Err(kigi_tool_runtime::ToolError::invalid_arguments(
|
||||
"image reference contained no data",
|
||||
));
|
||||
}
|
||||
|
||||
let (compressed, mime) = compress_reference(raw_bytes)?;
|
||||
let b64 = base64::engine::general_purpose::STANDARD.encode(&compressed);
|
||||
Ok(format!("data:{mime};base64,{b64}"))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Attachment reference resolution
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Parse an attached-image reference token into its 1-based display number.
|
||||
///
|
||||
/// Accepts the forms the model naturally produces for an image the user
|
||||
/// attached to the conversation: `[Image #1]`, `Image #1`, `image #1`, or
|
||||
/// a bare `#1`. Returns `None` for anything else — filesystem paths and
|
||||
/// `data:` / `file://` URLs fall through to direct resolution.
|
||||
fn parse_attachment_token(value: &str) -> Option<usize> {
|
||||
let trimmed = value.trim();
|
||||
// Strip optional surrounding brackets: `[…]`.
|
||||
let inner = trimmed
|
||||
.strip_prefix('[')
|
||||
.and_then(|s| s.strip_suffix(']'))
|
||||
.unwrap_or(trimmed)
|
||||
.trim();
|
||||
// Strip an optional leading `image` label (case-insensitive). The
|
||||
// 5-byte prefix is ASCII, so slicing at byte 5 stays on a boundary.
|
||||
let rest = match inner.get(..5).map(str::to_ascii_lowercase).as_deref() {
|
||||
Some("image") => inner[5..].trim_start(),
|
||||
_ => inner,
|
||||
};
|
||||
// Require the `#` sigil followed by a bare positive integer.
|
||||
let digits = rest.strip_prefix('#')?.trim();
|
||||
match digits.parse::<usize>() {
|
||||
Ok(n) if n >= 1 => Some(n),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve a single `image` argument to a reference `resolve_to_data_url`
|
||||
/// can read.
|
||||
///
|
||||
/// Attachment tokens (`[Image #N]`) are mapped to the durable reference
|
||||
/// the shell recorded for the current turn; everything else (filesystem
|
||||
/// paths, `data:` / `file://` URLs) passes through unchanged.
|
||||
fn resolve_attachment_reference(
|
||||
reference: &str,
|
||||
attached: Option<&crate::types::resources::AttachedImages>,
|
||||
) -> Result<String, kigi_tool_runtime::ToolError> {
|
||||
let Some(n) = parse_attachment_token(reference) else {
|
||||
return Ok(reference.to_owned());
|
||||
};
|
||||
let registry = attached.filter(|a| !a.0.is_empty()).ok_or_else(|| {
|
||||
// Tokens only resolve against the current message's attachments. An
|
||||
// empty registry usually means the image was attached in an earlier
|
||||
// message (cross-turn editing isn't supported yet), so steer the
|
||||
// model to ask for a re-attach rather than retry the dead token.
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"image reference {reference:?} matches no image attached to this message. If it was \
|
||||
attached earlier in the conversation, ask the user to re-attach it here; otherwise \
|
||||
pass an absolute filesystem path or a data: URL."
|
||||
))
|
||||
})?;
|
||||
registry.reference_for(n).map(str::to_owned).ok_or_else(|| {
|
||||
let available: Vec<String> = registry
|
||||
.0
|
||||
.iter()
|
||||
.map(|(num, _)| format!("[Image #{num}]"))
|
||||
.collect();
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"image reference {reference:?} does not match any attached image. Available: {}.",
|
||||
available.join(", ")
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tool input / schema
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
|
||||
pub struct ImageEditInput {
|
||||
#[schemars(
|
||||
description = "A text description of the desired edit or transformation. Describe what the output image should look like, referencing the input image(s)."
|
||||
)]
|
||||
pub prompt: String,
|
||||
|
||||
#[schemars(
|
||||
description = "Reference image(s) to condition the edit on. Each is one reference, in priority order: (1) a user attachment — its placeholder token, e.g. \"[Image #1]\" (attachments have no path you can see, so never invent one); (2) an absolute filesystem path the user gave you; (3) a `data:image/...;base64,...` URL."
|
||||
)]
|
||||
pub image: Vec<String>,
|
||||
|
||||
#[serde(default = "default_aspect_ratio")]
|
||||
#[schemars(
|
||||
description = "The aspect ratio of the output image. For single-image edits this is ignored — the output matches the input image's aspect ratio. For multi-image edits, defaults to 'auto'. Supported values: 1:1, 16:9, 9:16, 4:3, 3:4, 3:2, 2:3, 2:1, 1:2, 19.5:9, 9:19.5, 20:9, 9:20, auto."
|
||||
)]
|
||||
pub aspect_ratio: String,
|
||||
}
|
||||
|
||||
fn default_aspect_ratio() -> String {
|
||||
"auto".to_owned()
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tool implementation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct ImageEditTool;
|
||||
|
||||
impl crate::types::tool_metadata::ToolMetadata for ImageEditTool {
|
||||
fn kind(&self) -> ToolKind {
|
||||
ToolKind::ImageGen
|
||||
}
|
||||
|
||||
fn tool_namespace(&self) -> ToolNamespace {
|
||||
ToolNamespace::GrokBuild
|
||||
}
|
||||
|
||||
fn description_template(&self) -> &str {
|
||||
r##"Edit or transform existing image(s) via the xAI Imagine API; use instead of image_gen for image-to-image work (preserve likeness, transfer style, remix). Returns the saved image's absolute path. When telling the user where it was saved, refer to it by its short session-relative path (e.g. `images/1.jpg`) rather than the absolute path, so it renders as a clickable link that opens the image. Each required `image` is one reference — a user-attachment token (e.g. "[Image #1]"), an absolute filesystem path, or a `data:image/...;base64,...` URL (see the `image` parameter for the resolution order and details)."##
|
||||
}
|
||||
|
||||
fn requires_expr(&self) -> Expr<ToolRequirement> {
|
||||
Expr::True
|
||||
}
|
||||
}
|
||||
|
||||
impl kigi_tool_runtime::Tool for ImageEditTool {
|
||||
type Args = ImageEditInput;
|
||||
type Output = ToolOutput;
|
||||
|
||||
fn id(&self) -> kigi_tool_protocol::ToolId {
|
||||
kigi_tool_protocol::ToolId::new("image_edit").expect("valid tool id")
|
||||
}
|
||||
|
||||
fn description(
|
||||
&self,
|
||||
_ctx: &::kigi_tool_runtime::ListToolsContext,
|
||||
) -> kigi_tool_types::ToolDescription {
|
||||
kigi_tool_types::ToolDescription::new(
|
||||
"image_edit",
|
||||
crate::types::tool_metadata::ToolMetadata::description_template(self),
|
||||
)
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> kigi_tool_protocol::ToolCapabilities {
|
||||
kigi_tool_protocol::ToolCapabilities {
|
||||
is_read_only: false,
|
||||
tool_scope: Some(kigi_tool_protocol::ToolScope::Write),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(
|
||||
name = "tool.image_edit",
|
||||
skip_all,
|
||||
fields(prompt_len = input.prompt.len(), num_images = input.image.len(), aspect_ratio = %input.aspect_ratio)
|
||||
)]
|
||||
async fn run(
|
||||
&self,
|
||||
ctx: kigi_tool_runtime::ToolCallContext,
|
||||
input: ImageEditInput,
|
||||
) -> Result<ToolOutput, kigi_tool_runtime::ToolError> {
|
||||
use crate::types::tool_metadata::shared_resources;
|
||||
let resources = shared_resources(&ctx)?;
|
||||
|
||||
if input.image.is_empty() {
|
||||
return Err(kigi_tool_runtime::ToolError::invalid_arguments(
|
||||
"image_edit requires at least one reference image. \
|
||||
Use image_gen for text-only generation.",
|
||||
));
|
||||
}
|
||||
|
||||
let client = {
|
||||
let res = resources.lock().await;
|
||||
res.require::<ImageGenClient>()?.clone()
|
||||
};
|
||||
|
||||
// Free / X Basic users are zero-limited on Imagine server-side; return
|
||||
// the upsell prose instead of a doomed request (shares `image_gen`'s
|
||||
// message and short-circuits before resolving any attachments).
|
||||
if client.is_tier_restricted() {
|
||||
return Ok(ToolOutput::Text(
|
||||
super::image_gen::TIER_RESTRICTED_UPSELL.into(),
|
||||
));
|
||||
}
|
||||
|
||||
// Snapshot the per-turn attachment registry so `[Image #N]` tokens
|
||||
// resolve to the real attachment (see `resolve_attachment_reference`).
|
||||
let attached_images = {
|
||||
let res = resources.lock().await;
|
||||
res.get::<crate::types::resources::AttachedImages>()
|
||||
.cloned()
|
||||
};
|
||||
|
||||
// Resolve all references to compressed data URLs.
|
||||
let mut data_urls = Vec::with_capacity(input.image.len());
|
||||
for r in &input.image {
|
||||
let resolved = resolve_attachment_reference(r, attached_images.as_ref())?;
|
||||
data_urls.push(resolve_to_data_url(&resolved).await?);
|
||||
}
|
||||
tracing::info!(count = data_urls.len(), "resolved image references");
|
||||
|
||||
let base = client.base_url().trim_end_matches('/');
|
||||
let url = format!("{base}/images/edits");
|
||||
|
||||
let mut payload = serde_json::json!({
|
||||
"model": XAI_IMAGINE_MODEL,
|
||||
"prompt": input.prompt,
|
||||
"n": 1,
|
||||
"resolution": "1k",
|
||||
"response_format": "b64_json",
|
||||
});
|
||||
|
||||
// API: single ref → "image" object; multiple → "images" array.
|
||||
// For single-image edits the API auto-detects aspect ratio from the
|
||||
// input image and ignores the `aspect_ratio` field. Only send it
|
||||
// for multi-image edits where the API needs an explicit ratio.
|
||||
let mut imgs: Vec<serde_json::Value> = data_urls
|
||||
.iter()
|
||||
.map(|u| serde_json::json!({ "url": u }))
|
||||
.collect();
|
||||
if imgs.len() == 1 {
|
||||
payload["image"] = imgs.pop().unwrap();
|
||||
} else {
|
||||
payload["images"] = serde_json::Value::Array(imgs);
|
||||
payload["aspect_ratio"] = serde_json::json!(input.aspect_ratio);
|
||||
}
|
||||
|
||||
let sent_bearer = client.current_bearer().await;
|
||||
let mut req = client.http().post(&url).json(&payload);
|
||||
if let Some(ref key) = sent_bearer {
|
||||
req = req.header(AUTHORIZATION, format!("Bearer {key}"));
|
||||
}
|
||||
|
||||
let response = req.send().await.map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"Image edit API request failed: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
let status = response.status();
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED {
|
||||
client.record_401_attribution(ToolConsumer::ImageGen, sent_bearer.as_deref());
|
||||
}
|
||||
if !status.is_success() {
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
let truncated: String = body.chars().take(200).collect();
|
||||
tracing::warn!(http_status = %status, "Imagine edit API error: {truncated}");
|
||||
return Err(kigi_tool_runtime::ToolError::new(
|
||||
kigi_tool_runtime::ToolErrorKind::Custom,
|
||||
format!("Image edit failed with HTTP {status}: {truncated}"),
|
||||
)
|
||||
.with_details(serde_json::json!({"code": "http_failure", "status": status.as_u16()})));
|
||||
}
|
||||
|
||||
let body = response.text().await.map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"Failed to read image edit response body: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
let resp_json: ImageGenResponse = serde_json::from_str(&body).map_err(|e| {
|
||||
let preview: String = body.chars().take(500).collect();
|
||||
tracing::warn!("Imagine edit API returned unparseable body: {preview}");
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"Failed to parse image edit response: {e} — body preview: {preview}"
|
||||
))
|
||||
})?;
|
||||
|
||||
let b64_data = resp_json.b64_data().unwrap_or("");
|
||||
if b64_data.is_empty() {
|
||||
return Err(kigi_tool_runtime::ToolError::invalid_arguments(
|
||||
"Image edit returned no image data.",
|
||||
));
|
||||
}
|
||||
|
||||
let image_bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(b64_data)
|
||||
.map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"Failed to decode base64 image data: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
let session_folder = {
|
||||
let res = resources.lock().await;
|
||||
res.require::<SessionFolder>()?.0.clone()
|
||||
};
|
||||
|
||||
let absolute_path = client
|
||||
.writer()
|
||||
.save(&session_folder, &image_bytes, None)
|
||||
.await
|
||||
.map_err(|e| kigi_tool_runtime::ToolError::invalid_arguments(e.to_string()))?;
|
||||
|
||||
tracing::info!(
|
||||
path = %absolute_path.display(),
|
||||
bytes = image_bytes.len(),
|
||||
"edited image saved to disk"
|
||||
);
|
||||
|
||||
Ok(ToolOutput::ImageEdit(MediaGenOutput::new(absolute_path)))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::types::tool_metadata::test_ctx_with_call_id;
|
||||
|
||||
#[test]
|
||||
fn tool_name_and_description() {
|
||||
let tool = ImageEditTool;
|
||||
assert_eq!(kigi_tool_runtime::Tool::id(&tool).as_str(), "image_edit");
|
||||
let desc = crate::types::tool_metadata::ToolMetadata::description_template(&tool);
|
||||
assert!(desc.contains("Edit or transform"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn input_deserialization() {
|
||||
let input: ImageEditInput =
|
||||
serde_json::from_str(r#"{"prompt": "anime style", "image": ["/Users/me/photo.jpg"]}"#)
|
||||
.unwrap();
|
||||
assert_eq!(input.prompt, "anime style");
|
||||
assert_eq!(input.image, vec!["/Users/me/photo.jpg"]);
|
||||
assert_eq!(input.aspect_ratio, "auto");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn input_requires_image() {
|
||||
// image field is required by schema — empty array is a runtime check.
|
||||
let input: ImageEditInput =
|
||||
serde_json::from_str(r#"{"prompt": "test", "image": []}"#).unwrap();
|
||||
assert!(input.image.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_empty_image_array() {
|
||||
let tool = ImageEditTool;
|
||||
let resources = crate::types::resources::Resources::new();
|
||||
let result = kigi_tool_runtime::Tool::run(
|
||||
&tool,
|
||||
test_ctx_with_call_id(resources.into_shared(), "test-call"),
|
||||
ImageEditInput {
|
||||
prompt: "test".into(),
|
||||
image: vec![],
|
||||
aspect_ratio: "auto".into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let err = result.unwrap_err().to_string();
|
||||
assert!(err.contains("at least one reference image"), "got: {err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn errors_when_client_missing() {
|
||||
let tool = ImageEditTool;
|
||||
let resources = crate::types::resources::Resources::new();
|
||||
let result = kigi_tool_runtime::Tool::run(
|
||||
&tool,
|
||||
test_ctx_with_call_id(resources.into_shared(), "test-call"),
|
||||
ImageEditInput {
|
||||
prompt: "test".into(),
|
||||
image: vec!["/some/path.jpg".into()],
|
||||
aspect_ratio: "auto".into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let err = result.unwrap_err().to_string();
|
||||
assert!(err.contains("missing required resource"), "got: {err}");
|
||||
}
|
||||
|
||||
// ── compress_reference ───────────────────────────────────────────
|
||||
|
||||
fn tiny_jpeg() -> Vec<u8> {
|
||||
use image::{DynamicImage, RgbImage};
|
||||
let img = DynamicImage::ImageRgb8(RgbImage::new(2, 2));
|
||||
let mut buf = Vec::new();
|
||||
img.write_to(
|
||||
&mut std::io::Cursor::new(&mut buf),
|
||||
image::ImageFormat::Jpeg,
|
||||
)
|
||||
.unwrap();
|
||||
buf
|
||||
}
|
||||
|
||||
fn tiny_png() -> Vec<u8> {
|
||||
use image::{DynamicImage, RgbaImage};
|
||||
let img = DynamicImage::ImageRgba8(RgbaImage::new(2, 2));
|
||||
let mut buf = Vec::new();
|
||||
img.write_to(&mut std::io::Cursor::new(&mut buf), image::ImageFormat::Png)
|
||||
.unwrap();
|
||||
buf
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compress_small_jpeg_passthrough() {
|
||||
let jpeg = tiny_jpeg();
|
||||
let (out, mime) = compress_reference(jpeg.clone()).unwrap();
|
||||
assert_eq!(out, jpeg);
|
||||
assert_eq!(mime, "image/jpeg");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compress_small_png_passthrough() {
|
||||
let png = tiny_png();
|
||||
let (out, mime) = compress_reference(png.clone()).unwrap();
|
||||
assert_eq!(out, png);
|
||||
assert_eq!(mime, "image/png");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compress_oversized_shrinks() {
|
||||
use image::{DynamicImage, RgbImage};
|
||||
let mut img = RgbImage::new(1600, 1600);
|
||||
for (i, px) in img.pixels_mut().enumerate() {
|
||||
let v = (i * 37 + 13) as u8;
|
||||
*px = image::Rgb([v, v.wrapping_add(80), v.wrapping_add(160)]);
|
||||
}
|
||||
let mut buf = Vec::new();
|
||||
let enc = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buf, 100);
|
||||
DynamicImage::ImageRgb8(img)
|
||||
.write_with_encoder(enc)
|
||||
.unwrap();
|
||||
assert!(buf.len() > MAX_REF_RAW_BYTES);
|
||||
|
||||
let (out, mime) = compress_reference(buf).unwrap();
|
||||
assert!(out.len() <= MAX_REF_RAW_BYTES);
|
||||
assert!(mime == "image/jpeg" || mime == "image/png");
|
||||
}
|
||||
|
||||
// ── resolve_to_data_url ──────────────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_filesystem_path() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("test.jpg");
|
||||
std::fs::write(&path, tiny_jpeg()).unwrap();
|
||||
let url = resolve_to_data_url(path.to_str().unwrap()).await.unwrap();
|
||||
assert!(url.starts_with("data:image/jpeg;base64,"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_data_url_roundtrip() {
|
||||
let jpeg = tiny_jpeg();
|
||||
let b64 = base64::engine::general_purpose::STANDARD.encode(&jpeg);
|
||||
let input = format!("data:image/jpeg;base64,{b64}");
|
||||
let url = resolve_to_data_url(&input).await.unwrap();
|
||||
assert!(url.starts_with("data:image/jpeg;base64,"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_missing_file_errors() {
|
||||
assert!(resolve_to_data_url("/nonexistent/image.jpg").await.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_malformed_data_url_errors() {
|
||||
assert!(resolve_to_data_url("data:image/jpeg").await.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_file_uri_reads_path() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("test.jpg");
|
||||
std::fs::write(&path, tiny_jpeg()).unwrap();
|
||||
let uri = format!("file://{}", path.display());
|
||||
let url = resolve_to_data_url(&uri).await.unwrap();
|
||||
assert!(url.starts_with("data:image/jpeg;base64,"));
|
||||
}
|
||||
|
||||
// ── parse_attachment_token ───────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn parse_attachment_token_accepts_known_forms() {
|
||||
assert_eq!(parse_attachment_token("[Image #1]"), Some(1));
|
||||
assert_eq!(parse_attachment_token("Image #2"), Some(2));
|
||||
assert_eq!(parse_attachment_token("image #3"), Some(3));
|
||||
assert_eq!(parse_attachment_token("[image #4]"), Some(4));
|
||||
assert_eq!(parse_attachment_token("Image#5"), Some(5));
|
||||
assert_eq!(parse_attachment_token("#6"), Some(6));
|
||||
assert_eq!(parse_attachment_token(" [Image #7] "), Some(7));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_attachment_token_rejects_non_tokens() {
|
||||
assert_eq!(parse_attachment_token("/Users/me/photo.jpg"), None);
|
||||
assert_eq!(parse_attachment_token("data:image/png;base64,AAAA"), None);
|
||||
assert_eq!(parse_attachment_token("file:///tmp/x.png"), None);
|
||||
assert_eq!(parse_attachment_token("[Image #0]"), None);
|
||||
assert_eq!(parse_attachment_token("[Image #]"), None);
|
||||
assert_eq!(parse_attachment_token("Image one"), None);
|
||||
assert_eq!(parse_attachment_token(""), None);
|
||||
}
|
||||
|
||||
// ── resolve_attachment_reference ─────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn resolve_reference_passes_through_non_tokens() {
|
||||
let resolved = resolve_attachment_reference("/Users/me/photo.jpg", None).unwrap();
|
||||
assert_eq!(resolved, "/Users/me/photo.jpg");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_reference_maps_token_to_registry() {
|
||||
let attached = crate::types::resources::AttachedImages(vec![
|
||||
(1, "/tmp/a.png".to_owned()),
|
||||
(2, "/tmp/b.png".to_owned()),
|
||||
]);
|
||||
assert_eq!(
|
||||
resolve_attachment_reference("[Image #1]", Some(&attached)).unwrap(),
|
||||
"/tmp/a.png"
|
||||
);
|
||||
assert_eq!(
|
||||
resolve_attachment_reference("Image #2", Some(&attached)).unwrap(),
|
||||
"/tmp/b.png"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_reference_maps_by_number_not_position() {
|
||||
// After a mid-compose chip removal the surviving numbers are
|
||||
// non-contiguous (`#1`, `#3`). Resolution must key on the number,
|
||||
// not the list position, or `[Image #3]` would resolve to the wrong
|
||||
// file (or wrongly error).
|
||||
let attached = crate::types::resources::AttachedImages(vec![
|
||||
(1, "/tmp/first.png".to_owned()),
|
||||
(3, "/tmp/third.png".to_owned()),
|
||||
]);
|
||||
assert_eq!(
|
||||
resolve_attachment_reference("[Image #3]", Some(&attached)).unwrap(),
|
||||
"/tmp/third.png"
|
||||
);
|
||||
// `[Image #2]` was removed → no match.
|
||||
assert!(resolve_attachment_reference("[Image #2]", Some(&attached)).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_reference_token_without_registry_errors() {
|
||||
let err = resolve_attachment_reference("[Image #1]", None)
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
assert!(err.contains("re-attach"), "got: {err}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_reference_unmatched_number_errors() {
|
||||
let attached = crate::types::resources::AttachedImages(vec![(1, "/tmp/a.png".to_owned())]);
|
||||
let err = resolve_attachment_reference("[Image #2]", Some(&attached))
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
assert!(err.contains("does not match"), "got: {err}");
|
||||
assert!(err.contains("[Image #1]"), "should list available: {err}");
|
||||
}
|
||||
}
|
||||
@@ -1,585 +0,0 @@
|
||||
//! `image_gen` tool — generates images via the xAI Imagine API and saves
|
||||
//! them to the local filesystem so the model can reference them in code
|
||||
//! (e.g. `<img src="images/hero.jpg">`).
|
||||
//!
|
||||
//! Architecture follows the same pattern as `web_search`:
|
||||
//!
|
||||
//! - [`ImageGenConfig`] is built from session credentials by the host and
|
||||
//! injected into the tool registry.
|
||||
//! - When `Enabled`, an [`ImageGenClient`] is constructed once and injected
|
||||
//! into `Resources`. The tool reads it at runtime via `resources.require()`.
|
||||
//! - When `Disabled`, the tool is not registered so the model never sees it.
|
||||
//!
|
||||
//! The generated image is written to `<session_folder>/images/<n>.jpg`
|
||||
//! where `<n>` is a session-scoped counter (1, 2, 3, ... — 1 token each).
|
||||
//! The tool returns the absolute path so the model can copy or move the
|
||||
//! image into the project working directory when it needs a persistent asset.
|
||||
|
||||
use base64::Engine as _;
|
||||
use reqwest::header::{AUTHORIZATION, CONTENT_TYPE, HeaderValue};
|
||||
|
||||
use crate::attribution::{SharedAttributionCallback, ToolConsumer};
|
||||
use crate::types::SharedApiKeyProvider;
|
||||
|
||||
use crate::types::output::{MediaGenOutput, ToolOutput};
|
||||
use crate::types::requirements::{Expr, ToolRequirement};
|
||||
use crate::types::resources::SessionFolder;
|
||||
use crate::types::tool::{ToolKind, ToolNamespace};
|
||||
|
||||
/// Default Imagine model for `image_gen`. Used unless an explicit
|
||||
/// `model_override` is supplied via `ImageGenConfig::Enabled`.
|
||||
const XAI_IMAGINE_MODEL: &str = "grok-imagine-image-quality";
|
||||
// Some Imagine models (e.g. `grok-imagine-image`, selectable via `model_override`)
|
||||
// expand the prompt then generate, and the proxy buffers
|
||||
// the whole image before sending any bytes — so the client may receive nothing
|
||||
// for well over a minute. Keep these generous so a slow-but-progressing
|
||||
// generation isn't cut off.
|
||||
const IMAGE_GEN_TIMEOUT_SECS: u64 = 300;
|
||||
const IMAGE_GEN_READ_TIMEOUT_SECS: u64 = 240;
|
||||
const DEFAULT_IMAGE_DIR: &str = "images";
|
||||
|
||||
pub use kigi_tools_api::slash_commands::{
|
||||
IMAGE_GEN_TOOL_NAME, IMAGINE_COMMAND_NAME, imagine_instruction, imagine_usage_message,
|
||||
};
|
||||
|
||||
/// Prose returned to the model (as a normal, successful tool result) when a
|
||||
/// free / X Basic user calls `image_gen` or `image_edit`. The model relays it
|
||||
/// to the user. The deliberate `/imagine` slash command shows the richer
|
||||
/// SuperGrok upsell modal instead; this covers the natural-language path.
|
||||
pub(crate) const TIER_RESTRICTED_UPSELL: &str = "Image generation is a SuperGrok feature and isn't available on the free or X Basic tier. Let the user know they can unlock image and video generation by upgrading to SuperGrok: https://grok.com/supergrok?referrer=grok-build. Do not retry this tool.";
|
||||
|
||||
/// HTTP client for xAI Imagine API. Cloned per-request; shares `Arc` state.
|
||||
#[derive(Clone)]
|
||||
pub struct ImageGenClient {
|
||||
http: reqwest::Client,
|
||||
base_url: String,
|
||||
/// Imagine model slug used by `generate()`. Selected at construction
|
||||
/// from `ImageGenConfig::model_override` (falling back to
|
||||
/// [`XAI_IMAGINE_MODEL`]). `image_edit` uses its own model and is
|
||||
/// unaffected.
|
||||
model: String,
|
||||
writer: super::storage::SessionFileWriter,
|
||||
api_key_provider: Option<SharedApiKeyProvider>,
|
||||
/// Optional 401-attribution hook. Hosts wire this so a 401 from the
|
||||
/// Imagine API emits an `auth_401_attribution` event with
|
||||
/// `consumer == "ImageGen"` for unified auth-failure telemetry.
|
||||
attribution_callback: Option<SharedAttributionCallback>,
|
||||
/// When `true`, the user is on a tier the Imagine server zero-limits
|
||||
/// (free / X Basic). `image_gen` / `image_edit` short-circuit before any
|
||||
/// HTTP call and return the SuperGrok upsell prose instead. See
|
||||
/// [`ImageGenClient::is_tier_restricted`].
|
||||
tier_restricted: bool,
|
||||
}
|
||||
|
||||
impl ImageGenClient {
|
||||
pub fn new(
|
||||
config: &ImageGenConfig,
|
||||
api_key_provider: Option<SharedApiKeyProvider>,
|
||||
) -> Result<Self, kigi_tool_runtime::ToolError> {
|
||||
let ImageGenConfig::Enabled {
|
||||
api_key,
|
||||
base_url,
|
||||
extra_headers,
|
||||
model_override,
|
||||
tier_restricted,
|
||||
..
|
||||
} = config
|
||||
else {
|
||||
return Err(kigi_tool_runtime::ToolError::invalid_arguments(
|
||||
"Cannot create ImageGenClient from disabled config",
|
||||
));
|
||||
};
|
||||
let model = model_override
|
||||
.clone()
|
||||
.filter(|m| !m.trim().is_empty())
|
||||
.unwrap_or_else(|| XAI_IMAGINE_MODEL.to_owned());
|
||||
|
||||
let mut headers = reqwest::header::HeaderMap::new();
|
||||
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
|
||||
// Always bake the static api_key as the default Authorization header.
|
||||
// The dynamic provider overrides per-request; this is the fallback.
|
||||
headers.insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_str(&format!("Bearer {api_key}")).map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"Invalid API key for header: {e}"
|
||||
))
|
||||
})?,
|
||||
);
|
||||
|
||||
extra_headers.into_iter().try_for_each(|(key, value)| {
|
||||
let header_name =
|
||||
reqwest::header::HeaderName::from_bytes(key.as_bytes()).map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"Invalid header name '{key}': {e}"
|
||||
))
|
||||
})?;
|
||||
let header_value = HeaderValue::from_str(value).map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"Invalid header value for '{key}': {e}"
|
||||
))
|
||||
})?;
|
||||
headers.insert(header_name, header_value);
|
||||
Ok::<(), kigi_tool_runtime::ToolError>(())
|
||||
})?;
|
||||
|
||||
let http = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(IMAGE_GEN_TIMEOUT_SECS))
|
||||
.read_timeout(std::time::Duration::from_secs(IMAGE_GEN_READ_TIMEOUT_SECS))
|
||||
.default_headers(headers)
|
||||
.build()
|
||||
.map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"Failed to build HTTP client: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(Self {
|
||||
http,
|
||||
base_url: base_url.clone(),
|
||||
model,
|
||||
writer: super::storage::SessionFileWriter::new(DEFAULT_IMAGE_DIR, "jpg"),
|
||||
api_key_provider,
|
||||
attribution_callback: None,
|
||||
tier_restricted: *tier_restricted,
|
||||
})
|
||||
}
|
||||
|
||||
/// Whether the current user's tier (free / X Basic) is zero-limited on
|
||||
/// Imagine server-side. `image_gen` / `image_edit` use this to short-circuit
|
||||
/// with the SuperGrok upsell instead of issuing a doomed request.
|
||||
pub(crate) fn is_tier_restricted(&self) -> bool {
|
||||
self.tier_restricted
|
||||
}
|
||||
|
||||
/// Wire a 401-attribution callback into this client. Idempotent;
|
||||
/// safe to call before or after the first request. Builder-style
|
||||
/// so `new()` callers that don't care can ignore it.
|
||||
pub fn with_attribution_callback(
|
||||
mut self,
|
||||
callback: Option<SharedAttributionCallback>,
|
||||
) -> Self {
|
||||
self.attribution_callback = callback;
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) async fn current_bearer(&self) -> Option<String> {
|
||||
crate::types::api_key_provider::resolve_bearer(self.api_key_provider.as_ref()).await
|
||||
}
|
||||
|
||||
pub(crate) fn record_401_attribution(&self, consumer: ToolConsumer, sent_bearer: Option<&str>) {
|
||||
crate::attribution::emit_401(self.attribution_callback.as_ref(), consumer, sent_bearer);
|
||||
}
|
||||
|
||||
pub(crate) fn base_url(&self) -> &str {
|
||||
&self.base_url
|
||||
}
|
||||
|
||||
pub(crate) fn http(&self) -> &reqwest::Client {
|
||||
&self.http
|
||||
}
|
||||
|
||||
pub(crate) fn writer(&self) -> &super::storage::SessionFileWriter {
|
||||
&self.writer
|
||||
}
|
||||
|
||||
pub async fn generate(
|
||||
&self,
|
||||
prompt: &str,
|
||||
aspect_ratio: &str,
|
||||
) -> Result<Vec<u8>, kigi_tool_runtime::ToolError> {
|
||||
let url = format!("{}/images/generations", self.base_url.trim_end_matches('/'));
|
||||
|
||||
let payload = serde_json::json!({
|
||||
"model": self.model,
|
||||
"prompt": prompt,
|
||||
"n": 1,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"resolution": "1k",
|
||||
"response_format": "b64_json",
|
||||
});
|
||||
|
||||
// Capture the bearer once so the request and the 401-attribution
|
||||
// emit see the same value (even if the provider rotates between
|
||||
// the send and the response handling).
|
||||
let sent_bearer = self.current_bearer().await;
|
||||
let mut req = self.http.post(&url).json(&payload);
|
||||
if let Some(ref key) = sent_bearer {
|
||||
req = req.header(AUTHORIZATION, format!("Bearer {key}"));
|
||||
}
|
||||
|
||||
let response = req.send().await.map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"Image generation API request failed: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
let status = response.status();
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED {
|
||||
self.record_401_attribution(ToolConsumer::ImageGen, sent_bearer.as_deref());
|
||||
}
|
||||
if !status.is_success() {
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
let truncated: String = body.chars().take(200).collect();
|
||||
tracing::warn!(http_status = %status, "Imagine API error: {truncated}");
|
||||
return Err(kigi_tool_runtime::ToolError::new(
|
||||
kigi_tool_runtime::ToolErrorKind::Custom,
|
||||
format!("Image generation failed with HTTP {status}: {truncated}"),
|
||||
)
|
||||
.with_details(serde_json::json!({"code": "http_failure", "status": status.as_u16()})));
|
||||
}
|
||||
|
||||
let body = response.text().await.map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"Failed to read image generation response body: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
let resp_json: ImageGenResponse = serde_json::from_str(&body).map_err(|e| {
|
||||
let preview: String = body.chars().take(500).collect();
|
||||
tracing::warn!("Imagine API returned unparseable body: {preview}");
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"Failed to parse image generation response: {e} — body preview: {preview}"
|
||||
))
|
||||
})?;
|
||||
|
||||
let b64_data = resp_json.b64_data().unwrap_or("");
|
||||
|
||||
if b64_data.is_empty() {
|
||||
return Err(kigi_tool_runtime::ToolError::invalid_arguments(
|
||||
"Image generation returned no image data.",
|
||||
));
|
||||
}
|
||||
|
||||
base64::engine::general_purpose::STANDARD
|
||||
.decode(b64_data)
|
||||
.map_err(|e| {
|
||||
kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"Failed to decode base64 image data: {e}"
|
||||
))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// `Enabled` means credentials are present; each tool has its own gate.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub enum ImageGenConfig {
|
||||
#[default]
|
||||
Disabled,
|
||||
Enabled {
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
extra_headers: indexmap::IndexMap<String, String>,
|
||||
image_gen_enabled: bool,
|
||||
image_edit_enabled: bool,
|
||||
/// Optional Imagine model override for `image_gen`. When `Some(non-empty)`,
|
||||
/// `image_gen` calls that model instead of the default quality model
|
||||
/// ([`XAI_IMAGINE_MODEL`]). Driven by the remote
|
||||
/// `image_gen_model_override` config flag. `image_edit` is unaffected.
|
||||
model_override: Option<String>,
|
||||
/// `true` when the user is on a tier the Imagine server zero-limits
|
||||
/// (free / X Basic). The tools stay advertised to the model, but
|
||||
/// `image_gen` / `image_edit` short-circuit at call time with the
|
||||
/// SuperGrok upsell prose instead of a doomed request. Set by the
|
||||
/// host from the subscription tier; always `false` for team /
|
||||
/// API-key / workspace callers.
|
||||
tier_restricted: bool,
|
||||
},
|
||||
}
|
||||
|
||||
impl ImageGenConfig {
|
||||
/// Credentials present — required to construct any of the clients.
|
||||
pub fn has_credentials(&self) -> bool {
|
||||
matches!(self, Self::Enabled { .. })
|
||||
}
|
||||
|
||||
pub fn image_gen_enabled(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::Enabled {
|
||||
image_gen_enabled: true,
|
||||
..
|
||||
}
|
||||
)
|
||||
}
|
||||
|
||||
pub fn image_edit_enabled(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::Enabled {
|
||||
image_edit_enabled: true,
|
||||
..
|
||||
}
|
||||
)
|
||||
}
|
||||
|
||||
/// The configured `image_gen` model override, if any. `None` means the
|
||||
/// default quality model ([`XAI_IMAGINE_MODEL`]) is used.
|
||||
pub fn model_override(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Enabled { model_override, .. } => {
|
||||
model_override.as_deref().filter(|m| !m.trim().is_empty())
|
||||
}
|
||||
Self::Disabled => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
|
||||
pub struct ImageGenInput {
|
||||
#[schemars(description = "Text description of the image to generate.")]
|
||||
pub prompt: String,
|
||||
|
||||
#[serde(default = "default_aspect_ratio")]
|
||||
#[schemars(
|
||||
description = "Aspect ratio of the generated image, decide it based on the user's request. Defaults to 'auto'. 1:1 for square (icons, profiles), 16:9 for wide (landscapes, cinematic), 9:16 for tall (phone wallpapers, stories), 3:2 for horizontal photos, 2:3 for vertical (portraits, posters)."
|
||||
)]
|
||||
pub aspect_ratio: String,
|
||||
}
|
||||
|
||||
fn default_aspect_ratio() -> String {
|
||||
"auto".to_owned()
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
pub struct ImageGenResponse {
|
||||
#[serde(default)]
|
||||
data: Vec<ImageGenData>,
|
||||
}
|
||||
|
||||
impl ImageGenResponse {
|
||||
pub fn b64_data(&self) -> Option<&str> {
|
||||
self.data.first().and_then(|d| d.b64_json.as_deref())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct ImageGenData {
|
||||
b64_json: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct ImageGenTool;
|
||||
|
||||
impl crate::types::tool_metadata::ToolMetadata for ImageGenTool {
|
||||
fn kind(&self) -> ToolKind {
|
||||
ToolKind::ImageGen
|
||||
}
|
||||
|
||||
fn tool_namespace(&self) -> ToolNamespace {
|
||||
ToolNamespace::GrokBuild
|
||||
}
|
||||
|
||||
fn description_template(&self) -> &str {
|
||||
"Generate a new image from a text description using Imagine; returns the saved image's absolute path. When telling the user where it was saved, refer to it by its short session-relative path (e.g. `images/1.jpg`) rather than the absolute path, so it renders as a clickable link that opens the image. To produce multiple images, emit multiple tool calls with distinct prompts."
|
||||
}
|
||||
|
||||
fn requires_expr(&self) -> Expr<ToolRequirement> {
|
||||
Expr::True
|
||||
}
|
||||
}
|
||||
|
||||
impl kigi_tool_runtime::Tool for ImageGenTool {
|
||||
type Args = ImageGenInput;
|
||||
type Output = ToolOutput;
|
||||
|
||||
fn id(&self) -> kigi_tool_protocol::ToolId {
|
||||
kigi_tool_protocol::ToolId::new("image_gen").expect("valid tool id")
|
||||
}
|
||||
|
||||
fn description(
|
||||
&self,
|
||||
_ctx: &::kigi_tool_runtime::ListToolsContext,
|
||||
) -> kigi_tool_types::ToolDescription {
|
||||
kigi_tool_types::ToolDescription::new(
|
||||
"image_gen",
|
||||
crate::types::tool_metadata::ToolMetadata::description_template(self),
|
||||
)
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> kigi_tool_protocol::ToolCapabilities {
|
||||
kigi_tool_protocol::ToolCapabilities {
|
||||
is_read_only: false,
|
||||
tool_scope: Some(kigi_tool_protocol::ToolScope::Write),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(
|
||||
name = "tool.image_gen",
|
||||
skip_all,
|
||||
fields(prompt_len = input.prompt.len(), aspect_ratio = %input.aspect_ratio)
|
||||
)]
|
||||
async fn run(
|
||||
&self,
|
||||
ctx: kigi_tool_runtime::ToolCallContext,
|
||||
input: ImageGenInput,
|
||||
) -> Result<ToolOutput, kigi_tool_runtime::ToolError> {
|
||||
use crate::types::tool_metadata::shared_resources;
|
||||
let resources = shared_resources(&ctx)?;
|
||||
|
||||
let client = {
|
||||
let res = resources.lock().await;
|
||||
res.require::<ImageGenClient>()?.clone()
|
||||
};
|
||||
|
||||
// Free / X Basic users are zero-limited on Imagine server-side; return
|
||||
// the upsell prose instead of a doomed request (the tool stays
|
||||
// advertised so the model can surface the nudge in-conversation).
|
||||
if client.is_tier_restricted() {
|
||||
return Ok(ToolOutput::Text(TIER_RESTRICTED_UPSELL.into()));
|
||||
}
|
||||
|
||||
let image_bytes = client.generate(&input.prompt, &input.aspect_ratio).await?;
|
||||
|
||||
let session_folder = {
|
||||
let res = resources.lock().await;
|
||||
res.require::<SessionFolder>()?.0.clone()
|
||||
};
|
||||
|
||||
let absolute_path = client
|
||||
.writer
|
||||
.save(&session_folder, &image_bytes, None)
|
||||
.await
|
||||
.map_err(|e| kigi_tool_runtime::ToolError::invalid_arguments(e.to_string()))?;
|
||||
|
||||
tracing::info!(
|
||||
path = %absolute_path.display(),
|
||||
bytes = image_bytes.len(),
|
||||
"image saved to disk"
|
||||
);
|
||||
|
||||
Ok(ToolOutput::ImageGen(MediaGenOutput::new(absolute_path)))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::types::tool_metadata::test_ctx_with_call_id;
|
||||
|
||||
#[test]
|
||||
fn tool_name_and_description() {
|
||||
let tool = ImageGenTool;
|
||||
assert_eq!(kigi_tool_runtime::Tool::id(&tool).as_str(), "image_gen");
|
||||
assert!(
|
||||
crate::types::tool_metadata::ToolMetadata::description_template(&tool)
|
||||
.contains("Generate a new image from a text description")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_aspect_ratio_is_auto() {
|
||||
let input: ImageGenInput = serde_json::from_str(r#"{"prompt": "test"}"#).unwrap();
|
||||
assert_eq!(input.aspect_ratio, "auto");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn per_tool_gates_are_independent() {
|
||||
let cfg = ImageGenConfig::Enabled {
|
||||
api_key: "k".into(),
|
||||
base_url: "https://api.x.ai/v1".into(),
|
||||
extra_headers: indexmap::IndexMap::new(),
|
||||
image_gen_enabled: false,
|
||||
image_edit_enabled: true,
|
||||
model_override: Some("grok-imagine-image".into()),
|
||||
tier_restricted: false,
|
||||
};
|
||||
assert!(cfg.has_credentials());
|
||||
assert!(!cfg.image_gen_enabled());
|
||||
assert!(cfg.image_edit_enabled());
|
||||
assert_eq!(cfg.model_override(), Some("grok-imagine-image"));
|
||||
|
||||
assert!(!ImageGenConfig::Disabled.has_credentials());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_selects_model_from_override() {
|
||||
let mk = |model_override: Option<&str>| ImageGenConfig::Enabled {
|
||||
api_key: "k".into(),
|
||||
base_url: "https://api.x.ai/v1".into(),
|
||||
extra_headers: indexmap::IndexMap::new(),
|
||||
image_gen_enabled: true,
|
||||
image_edit_enabled: true,
|
||||
model_override: model_override.map(String::from),
|
||||
tier_restricted: false,
|
||||
};
|
||||
// No override → default quality model.
|
||||
assert_eq!(
|
||||
ImageGenClient::new(&mk(None), None).unwrap().model,
|
||||
XAI_IMAGINE_MODEL
|
||||
);
|
||||
// Empty override → treated as no override.
|
||||
assert_eq!(
|
||||
ImageGenClient::new(&mk(Some("")), None).unwrap().model,
|
||||
XAI_IMAGINE_MODEL
|
||||
);
|
||||
// Override → that exact model slug.
|
||||
assert_eq!(
|
||||
ImageGenClient::new(&mk(Some("grok-imagine-image")), None)
|
||||
.unwrap()
|
||||
.model,
|
||||
"grok-imagine-image"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn errors_when_client_missing() {
|
||||
let tool = ImageGenTool;
|
||||
let resources = crate::types::resources::Resources::new();
|
||||
let result = kigi_tool_runtime::Tool::run(
|
||||
&tool,
|
||||
test_ctx_with_call_id(resources.into_shared(), "test-call"),
|
||||
ImageGenInput {
|
||||
prompt: "a test image".into(),
|
||||
aspect_ratio: "auto".into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(result.is_err());
|
||||
let err_msg = result.unwrap_err().to_string();
|
||||
assert!(
|
||||
err_msg.contains("missing required resource"),
|
||||
"Expected MissingResource error, got: {err_msg}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tier_restricted_short_circuits_with_upsell() {
|
||||
// A free / X Basic user's image_gen call returns the SuperGrok upsell
|
||||
// prose as a normal result (no HTTP, no error card) so the model can
|
||||
// relay it. Only the client is inserted — the short-circuit returns
|
||||
// before any other resource (e.g. SessionFolder) is required.
|
||||
let cfg = ImageGenConfig::Enabled {
|
||||
api_key: "k".into(),
|
||||
base_url: "https://api.x.ai/v1".into(),
|
||||
extra_headers: indexmap::IndexMap::new(),
|
||||
image_gen_enabled: true,
|
||||
image_edit_enabled: true,
|
||||
model_override: None,
|
||||
tier_restricted: true,
|
||||
};
|
||||
let mut resources = crate::types::resources::Resources::new();
|
||||
resources.insert(ImageGenClient::new(&cfg, None).unwrap());
|
||||
|
||||
let result = kigi_tool_runtime::Tool::run(
|
||||
&ImageGenTool,
|
||||
test_ctx_with_call_id(resources.into_shared(), "test-call"),
|
||||
ImageGenInput {
|
||||
prompt: "a cat".into(),
|
||||
aspect_ratio: "auto".into(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("tier-restricted call must succeed with upsell prose");
|
||||
|
||||
match result {
|
||||
ToolOutput::Text(t) => {
|
||||
assert!(t.text.contains("SuperGrok"), "got: {}", t.text);
|
||||
assert!(t.text.contains("supergrok?referrer=grok-build"));
|
||||
}
|
||||
other => panic!("expected Text upsell, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -15,8 +15,6 @@ pub mod deploy_app;
|
||||
pub mod enter_plan_mode;
|
||||
pub mod exit_plan_mode;
|
||||
pub mod grep;
|
||||
pub mod image_edit;
|
||||
pub mod image_gen;
|
||||
pub mod kill_task;
|
||||
pub mod list_dir;
|
||||
pub mod lsp;
|
||||
@@ -29,7 +27,6 @@ pub mod task;
|
||||
pub mod task_output;
|
||||
pub mod todo;
|
||||
pub mod update_goal;
|
||||
pub mod video_gen;
|
||||
pub mod web_fetch;
|
||||
pub mod web_search;
|
||||
pub use ask_user_question::AskUserQuestionTool;
|
||||
@@ -38,11 +35,6 @@ pub use deploy_app::{AppBuilderDeployerConfig, DEPLOY_APP_TOOL_NAME};
|
||||
pub use enter_plan_mode::EnterPlanModeTool;
|
||||
pub use exit_plan_mode::ExitPlanModeTool;
|
||||
pub use grep::GrepTool;
|
||||
pub use image_edit::{IMAGE_EDIT_TOOL_NAME, ImageEditTool};
|
||||
pub use image_gen::{
|
||||
IMAGE_GEN_TOOL_NAME, IMAGINE_COMMAND_NAME, ImageGenTool, imagine_instruction,
|
||||
imagine_usage_message,
|
||||
};
|
||||
pub use kill_task::{KillTaskTool, KillTerminalCommandTool};
|
||||
pub use list_dir::ListDirTool;
|
||||
pub use lsp::LspTool;
|
||||
@@ -58,10 +50,5 @@ pub use task::TaskTool;
|
||||
pub use task_output::{GetTerminalCommandOutputTool, TaskOutputTool, WaitTasksTool};
|
||||
pub use todo::TodoWriteTool;
|
||||
pub use update_goal::{UPDATE_GOAL_TOOL_NAME, UpdateGoalTool};
|
||||
pub use video_gen::{
|
||||
IMAGE_TO_VIDEO_TOOL_NAME, IMAGINE_VIDEO_COMMAND_NAME, ImageToVideoTool,
|
||||
REFERENCE_TO_VIDEO_TOOL_NAME, ReferenceToVideoTool, imagine_video_instruction,
|
||||
imagine_video_usage_message,
|
||||
};
|
||||
pub use web_fetch::{WebFetchClient, WebFetchConfig, WebFetchParams, WebFetchTool};
|
||||
pub use web_search::WebSearchTool;
|
||||
|
||||
@@ -237,10 +237,6 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode {
|
||||
ToolKind::MemoryGet,
|
||||
ToolKind::WebSearch,
|
||||
ToolKind::WebFetch,
|
||||
ToolKind::ImageGen,
|
||||
ToolKind::VideoGen,
|
||||
ToolKind::ImageToVideo,
|
||||
ToolKind::ReferenceToVideo,
|
||||
ToolKind::BackgroundTaskAction,
|
||||
ToolKind::KillTaskAction,
|
||||
ToolKind::Task,
|
||||
@@ -285,10 +281,6 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode {
|
||||
ToolKind::MemoryGet,
|
||||
ToolKind::WebSearch,
|
||||
ToolKind::WebFetch,
|
||||
ToolKind::ImageGen,
|
||||
ToolKind::VideoGen,
|
||||
ToolKind::ImageToVideo,
|
||||
ToolKind::ReferenceToVideo,
|
||||
ToolKind::BackgroundTaskAction,
|
||||
ToolKind::KillTaskAction,
|
||||
ToolKind::Task,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -3,7 +3,7 @@
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::types::output::{MCPOutput, ToolOutput};
|
||||
use crate::types::output::ToolOutput;
|
||||
use crate::types::tool::{ToolKind, ToolNamespace};
|
||||
use crate::util::mcp_truncate::{McpTruncateContext, truncate_tool_output};
|
||||
|
||||
@@ -94,48 +94,6 @@ async fn dispatch_local_mcp(
|
||||
.map_err(|e| kigi_tool_runtime::ToolError::custom("output_decoding", e.to_string()))
|
||||
}
|
||||
|
||||
fn gateway_result_is_error(result: &serde_json::Value) -> bool {
|
||||
result
|
||||
.get("isError")
|
||||
.or_else(|| result.get("is_error"))
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn gateway_result_to_text(result: serde_json::Value) -> String {
|
||||
if let Some(content) = result.get("content").and_then(|v| v.as_array()) {
|
||||
let parts: Vec<String> = content
|
||||
.iter()
|
||||
.filter_map(|item| {
|
||||
if item.get("type").and_then(|v| v.as_str()) == Some("text") {
|
||||
item.get("text").and_then(|v| v.as_str()).map(str::to_owned)
|
||||
} else if item.get("type").and_then(|v| v.as_str()) == Some("image") {
|
||||
let mime = item
|
||||
.get("mimeType")
|
||||
.or_else(|| item.get("mime_type"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("image/png");
|
||||
item.get("data")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|data| format!("data:{mime};base64,{data}"))
|
||||
} else if item.get("type").and_then(|v| v.as_str()) == Some("resource") {
|
||||
serde_json::to_string(item).ok()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
if !parts.is_empty() {
|
||||
return parts.join("\n");
|
||||
}
|
||||
}
|
||||
|
||||
match result {
|
||||
serde_json::Value::String(s) => s,
|
||||
other => serde_json::to_string_pretty(&other).unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_mcp_arguments(input: serde_json::Value) -> serde_json::Value {
|
||||
match input {
|
||||
serde_json::Value::String(s) => match serde_json::from_str(&s) {
|
||||
@@ -147,54 +105,6 @@ fn normalize_mcp_arguments(input: serde_json::Value) -> serde_json::Value {
|
||||
}
|
||||
}
|
||||
|
||||
fn is_local_tool_id_rejection(err: &kigi_tool_runtime::ToolError, tool_name: &str) -> bool {
|
||||
err.kind == kigi_tool_runtime::ToolErrorKind::InvalidArguments
|
||||
&& err.detail == format!("invalid tool name: '{tool_name}'")
|
||||
}
|
||||
|
||||
async fn gateway_lookup(
|
||||
ctx: &kigi_tool_runtime::ToolCallContext,
|
||||
tool_name: &str,
|
||||
) -> (
|
||||
Option<crate::types::resources::ManagedGatewayToolSource>,
|
||||
Option<crate::types::resources::ManagedGatewayToolClient>,
|
||||
) {
|
||||
let Some(resources) = crate::types::tool_metadata::shared_resources(ctx).ok() else {
|
||||
return (None, None);
|
||||
};
|
||||
let guard = resources.lock().await;
|
||||
let source = guard
|
||||
.get::<crate::types::resources::ManagedGatewayToolCatalog>()
|
||||
.and_then(|catalog| catalog.get(tool_name).cloned());
|
||||
let client = guard
|
||||
.get::<crate::types::resources::ManagedGatewayToolClient>()
|
||||
.cloned()
|
||||
.filter(|_| source.is_some());
|
||||
(source, client)
|
||||
}
|
||||
|
||||
fn gateway_response_to_output(
|
||||
tool_name: &str,
|
||||
source: crate::types::resources::ManagedGatewayToolSource,
|
||||
result: serde_json::Value,
|
||||
) -> ToolOutput {
|
||||
let is_error = gateway_result_is_error(&result);
|
||||
let text = gateway_result_to_text(result);
|
||||
if is_error {
|
||||
ToolOutput::MCP(MCPOutput::errored(
|
||||
tool_name.to_owned(),
|
||||
source.connector_name,
|
||||
text,
|
||||
))
|
||||
} else {
|
||||
ToolOutput::MCP(MCPOutput::okay_output(
|
||||
tool_name.to_owned(),
|
||||
source.connector_name,
|
||||
text,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn dispatch_mcp_tool(
|
||||
ctx: &kigi_tool_runtime::ToolCallContext,
|
||||
tool_name: &str,
|
||||
@@ -202,71 +112,16 @@ pub async fn dispatch_mcp_tool(
|
||||
caller: &str,
|
||||
) -> Result<ToolOutput, kigi_tool_runtime::ToolError> {
|
||||
let tool_input = normalize_mcp_arguments(tool_input);
|
||||
let (gateway_source, gateway_client) = gateway_lookup(ctx, tool_name).await;
|
||||
let dispatch = ctx
|
||||
let Some(dispatch) = ctx
|
||||
.extensions
|
||||
.get::<crate::types::resources::InnerDispatch>();
|
||||
|
||||
if gateway_source.is_none() && dispatch.is_none() {
|
||||
.get::<crate::types::resources::InnerDispatch>()
|
||||
else {
|
||||
return Err(kigi_tool_runtime::ToolError::invalid_arguments(format!(
|
||||
"{caller} called outside of tool execution context. inner_dispatch not set -- this is a bug."
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(source) = gateway_source {
|
||||
// A gateway-catalog name can collide with a local `server__tool` MCP
|
||||
// tool. Local wins on a name clash: probe local dispatch first and only
|
||||
// fall through to the gateway when the local side reports the tool as
|
||||
// not found, or rejects the catalog-derived name as an invalid local
|
||||
// ToolId. A real error from a local tool that actually dispatched
|
||||
// propagates instead of silently retrying against the gateway.
|
||||
if tool_name.contains("__")
|
||||
&& let Some(dispatch) = dispatch.clone()
|
||||
{
|
||||
match dispatch_local_mcp(dispatch, tool_name, tool_input.clone(), ctx.clone()).await {
|
||||
Ok(local_output) => return Ok(local_output),
|
||||
Err(err)
|
||||
if err.kind != kigi_tool_runtime::ToolErrorKind::NotFound
|
||||
&& !is_local_tool_id_rejection(&err, tool_name) =>
|
||||
{
|
||||
return Err(err);
|
||||
}
|
||||
Err(_) => {}
|
||||
}
|
||||
}
|
||||
|
||||
let Some(client) = gateway_client else {
|
||||
return Err(kigi_tool_runtime::ToolError::custom(
|
||||
"managed_gateway_unavailable",
|
||||
format!(
|
||||
"Managed MCP gateway tool '{}' is indexed but no gateway client is available.",
|
||||
tool_name
|
||||
),
|
||||
));
|
||||
};
|
||||
let response = client
|
||||
.0
|
||||
.call_tool(&source.call_id, tool_input, caller)
|
||||
.await?;
|
||||
tracing::debug!(
|
||||
tool_name = %tool_name,
|
||||
reauth = response.connectors_needing_reauth.len(),
|
||||
"Managed MCP gateway tool call completed"
|
||||
);
|
||||
return Ok(gateway_response_to_output(
|
||||
tool_name,
|
||||
source,
|
||||
response.result,
|
||||
));
|
||||
}
|
||||
|
||||
dispatch_local_mcp(
|
||||
dispatch.expect("dispatch is set for local MCP path"),
|
||||
tool_name,
|
||||
tool_input,
|
||||
ctx.clone(),
|
||||
)
|
||||
.await
|
||||
dispatch_local_mcp(dispatch, tool_name, tool_input, ctx.clone()).await
|
||||
}
|
||||
|
||||
impl crate::types::tool_metadata::ToolMetadata for UseTool {
|
||||
@@ -316,33 +171,29 @@ impl kigi_tool_runtime::Tool for UseTool {
|
||||
ctx: kigi_tool_runtime::ToolCallContext,
|
||||
input: UseToolInput,
|
||||
) -> Result<ToolOutput, kigi_tool_runtime::ToolError> {
|
||||
use crate::types::resources::{EnabledNativeToolNames, ManagedGatewayToolCatalog, Params};
|
||||
use crate::types::resources::{EnabledNativeToolNames, Params};
|
||||
|
||||
let resources = crate::types::tool_metadata::shared_resources(&ctx).ok();
|
||||
let (gateway_source, is_native, search_tool_name) =
|
||||
if let Some(resources) = resources.as_ref() {
|
||||
let guard = resources.lock().await;
|
||||
let gateway_source = guard
|
||||
.get::<ManagedGatewayToolCatalog>()
|
||||
.and_then(|catalog| catalog.get(&input.tool_name).cloned());
|
||||
let correction_enabled = guard
|
||||
.get::<Params<UseToolParams>>()
|
||||
.is_none_or(|p| p.0.native_tool_correction);
|
||||
let native = correction_enabled
|
||||
&& guard
|
||||
.get::<EnabledNativeToolNames>()
|
||||
.is_some_and(|set| set.contains(&input.tool_name));
|
||||
let st = guard
|
||||
.get::<crate::types::template_renderer::TemplateRenderer>()
|
||||
.and_then(|r| r.tool_for_kind(ToolKind::SearchTool))
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(|| "search_tool".to_string());
|
||||
(gateway_source, native, st)
|
||||
} else {
|
||||
(None, false, "search_tool".to_string())
|
||||
};
|
||||
let (is_native, search_tool_name) = if let Some(resources) = resources.as_ref() {
|
||||
let guard = resources.lock().await;
|
||||
let correction_enabled = guard
|
||||
.get::<Params<UseToolParams>>()
|
||||
.is_none_or(|p| p.0.native_tool_correction);
|
||||
let native = correction_enabled
|
||||
&& guard
|
||||
.get::<EnabledNativeToolNames>()
|
||||
.is_some_and(|set| set.contains(&input.tool_name));
|
||||
let st = guard
|
||||
.get::<crate::types::template_renderer::TemplateRenderer>()
|
||||
.and_then(|r| r.tool_for_kind(ToolKind::SearchTool))
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(|| "search_tool".to_string());
|
||||
(native, st)
|
||||
} else {
|
||||
(false, "search_tool".to_string())
|
||||
};
|
||||
|
||||
if !input.tool_name.contains("__") && gateway_source.is_none() {
|
||||
if !input.tool_name.contains("__") {
|
||||
return Err(if is_native {
|
||||
// Native tool wrongly routed through use_tool. Tell the model
|
||||
// to call it directly. Strategy chosen via offline eval over
|
||||
@@ -448,39 +299,6 @@ mod tests {
|
||||
(ctx, args)
|
||||
}
|
||||
|
||||
struct NotFoundDispatch;
|
||||
|
||||
struct InvalidArgumentsDispatch;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl kigi_tool_runtime::ToolDispatch for NotFoundDispatch {
|
||||
async fn call(
|
||||
&self,
|
||||
tool_id: kigi_tool_protocol::ToolId,
|
||||
_args: serde_json::Value,
|
||||
_ctx: kigi_tool_runtime::ToolCallContext,
|
||||
) -> kigi_tool_runtime::ToolStream<kigi_tool_runtime::TypedToolOutput> {
|
||||
kigi_tool_runtime::terminal_only(Err(kigi_tool_runtime::ToolError::not_found(
|
||||
tool_id,
|
||||
"Tool not found",
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl kigi_tool_runtime::ToolDispatch for InvalidArgumentsDispatch {
|
||||
async fn call(
|
||||
&self,
|
||||
_tool_id: kigi_tool_protocol::ToolId,
|
||||
_args: serde_json::Value,
|
||||
_ctx: kigi_tool_runtime::ToolCallContext,
|
||||
) -> kigi_tool_runtime::ToolStream<kigi_tool_runtime::TypedToolOutput> {
|
||||
kigi_tool_runtime::terminal_only(Err(kigi_tool_runtime::ToolError::invalid_arguments(
|
||||
"local validation failed",
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
/// Mock dispatch that always returns an error.
|
||||
struct ErrorToolDispatch {
|
||||
error: String,
|
||||
@@ -604,238 +422,6 @@ mod tests {
|
||||
assert!(err.detail.contains("bad__tool"));
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct MockGatewayCaller {
|
||||
captured: SharedArgs,
|
||||
result: serde_json::Value,
|
||||
expected_call_id: Option<&'static str>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl crate::types::resources::ManagedGatewayToolCaller for MockGatewayCaller {
|
||||
async fn call_tool(
|
||||
&self,
|
||||
call_id: &str,
|
||||
arguments: serde_json::Value,
|
||||
_caller: &str,
|
||||
) -> Result<
|
||||
crate::types::resources::ManagedGatewayToolCallResponse,
|
||||
kigi_tool_runtime::ToolError,
|
||||
> {
|
||||
if let Some(expected) = self.expected_call_id {
|
||||
assert_eq!(call_id, expected);
|
||||
}
|
||||
*self.captured.lock().unwrap() = Some(arguments);
|
||||
Ok(crate::types::resources::ManagedGatewayToolCallResponse {
|
||||
result: self.result.clone(),
|
||||
connectors_needing_reauth: vec![],
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn gateway_resources(
|
||||
captured: SharedArgs,
|
||||
result: serde_json::Value,
|
||||
) -> crate::types::resources::SharedResources {
|
||||
gateway_resources_with_expected_call_id(captured, result, Some("grafana.searchDashboards"))
|
||||
}
|
||||
|
||||
fn gateway_resources_with_expected_call_id(
|
||||
captured: SharedArgs,
|
||||
result: serde_json::Value,
|
||||
expected_call_id: Option<&'static str>,
|
||||
) -> crate::types::resources::SharedResources {
|
||||
use crate::types::resources::{
|
||||
ManagedGatewayToolCatalog, ManagedGatewayToolClient, ManagedGatewayToolSource,
|
||||
Resources,
|
||||
};
|
||||
let mut resources = Resources::new();
|
||||
resources.insert(ManagedGatewayToolCatalog(std::collections::HashMap::from(
|
||||
[
|
||||
(
|
||||
"grafana__search_dashboards".to_string(),
|
||||
ManagedGatewayToolSource {
|
||||
connector_id: "grafana".to_string(),
|
||||
connector_name: "Grafana".to_string(),
|
||||
tool_id: "search_dashboards".to_string(),
|
||||
tool_name: "Search Dashboards".to_string(),
|
||||
call_id: "grafana.searchDashboards".to_string(),
|
||||
},
|
||||
),
|
||||
(
|
||||
"server__tool".to_string(),
|
||||
ManagedGatewayToolSource {
|
||||
connector_id: "server".to_string(),
|
||||
connector_name: "Gateway Collision".to_string(),
|
||||
tool_id: "tool".to_string(),
|
||||
tool_name: "Tool".to_string(),
|
||||
call_id: "gateway.collision".to_string(),
|
||||
},
|
||||
),
|
||||
(
|
||||
"connector__bad/id".to_string(),
|
||||
ManagedGatewayToolSource {
|
||||
connector_id: "connector".to_string(),
|
||||
connector_name: "Gateway Invalid Local".to_string(),
|
||||
tool_id: "bad/id".to_string(),
|
||||
tool_name: "Bad ID".to_string(),
|
||||
call_id: "gateway.invalidLocal".to_string(),
|
||||
},
|
||||
),
|
||||
],
|
||||
)));
|
||||
resources.insert(ManagedGatewayToolClient(Arc::new(MockGatewayCaller {
|
||||
captured,
|
||||
result,
|
||||
expected_call_id,
|
||||
})));
|
||||
resources.into_shared()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_tool_dispatches_to_gateway_call_id() {
|
||||
let captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
|
||||
let ctx = ctx_with_dispatch_and_resources(
|
||||
NotFoundDispatch,
|
||||
gateway_resources(
|
||||
Arc::clone(&captured),
|
||||
serde_json::json!({"content": [{"type": "text", "text": "dashboards"}]}),
|
||||
),
|
||||
);
|
||||
|
||||
let result = kigi_tool_runtime::Tool::run(
|
||||
&UseTool,
|
||||
ctx,
|
||||
UseToolInput {
|
||||
tool_name: "grafana__search_dashboards".into(),
|
||||
tool_input: serde_json::json!({"query": "prod"}),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(captured.lock().unwrap().clone().unwrap()["query"], "prod");
|
||||
if let ToolOutput::MCP(mcp) = result {
|
||||
match mcp.output() {
|
||||
crate::types::output::MCPOutputDetails::OkayOutput(text) => {
|
||||
assert_eq!(text, "dashboards")
|
||||
}
|
||||
_ => panic!("expected okay output"),
|
||||
}
|
||||
} else {
|
||||
panic!("expected gateway result to map to MCP output");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_error_result_maps_to_mcp_error() {
|
||||
let captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
|
||||
let ctx = ctx_with_dispatch_and_resources(
|
||||
NotFoundDispatch,
|
||||
gateway_resources(
|
||||
Arc::clone(&captured),
|
||||
serde_json::json!({
|
||||
"isError": true,
|
||||
"content": [{"type": "text", "text": "remote failed"}]
|
||||
}),
|
||||
),
|
||||
);
|
||||
|
||||
let result = kigi_tool_runtime::Tool::run(
|
||||
&UseTool,
|
||||
ctx,
|
||||
UseToolInput {
|
||||
tool_name: "grafana__search_dashboards".into(),
|
||||
tool_input: serde_json::json!({}),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_error());
|
||||
assert!(
|
||||
result
|
||||
.to_prompt_format()
|
||||
.contains("Failed to call grafana__search_dashboards: remote failed")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_snake_case_error_result_maps_to_mcp_error() {
|
||||
let captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
|
||||
let ctx = ctx_with_dispatch_and_resources(
|
||||
NotFoundDispatch,
|
||||
gateway_resources(
|
||||
Arc::clone(&captured),
|
||||
serde_json::json!({
|
||||
"is_error": true,
|
||||
"content": [{"type": "text", "text": "snake failed"}]
|
||||
}),
|
||||
),
|
||||
);
|
||||
|
||||
let result = kigi_tool_runtime::Tool::run(
|
||||
&UseTool,
|
||||
ctx,
|
||||
UseToolInput {
|
||||
tool_name: "grafana__search_dashboards".into(),
|
||||
tool_input: serde_json::json!({}),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_error());
|
||||
assert!(result.to_prompt_format().contains("snake failed"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_call_result_converts_to_model_visible_output() {
|
||||
let captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
|
||||
let ctx = ctx_with_dispatch_and_resources(
|
||||
NotFoundDispatch,
|
||||
gateway_resources(Arc::clone(&captured), serde_json::json!({"ok": true})),
|
||||
);
|
||||
|
||||
let result = kigi_tool_runtime::Tool::run(
|
||||
&UseTool,
|
||||
ctx,
|
||||
UseToolInput {
|
||||
tool_name: "grafana__search_dashboards".into(),
|
||||
tool_input: serde_json::json!({}),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.to_prompt_format().contains("\"ok\": true"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_null_arguments_default_to_object() {
|
||||
let captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
|
||||
let ctx = ctx_with_dispatch_and_resources(
|
||||
NotFoundDispatch,
|
||||
gateway_resources(Arc::clone(&captured), serde_json::json!("ok")),
|
||||
);
|
||||
|
||||
kigi_tool_runtime::Tool::run(
|
||||
&UseTool,
|
||||
ctx,
|
||||
UseToolInput {
|
||||
tool_name: "grafana__search_dashboards".into(),
|
||||
tool_input: serde_json::Value::Null,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
captured.lock().unwrap().clone().unwrap(),
|
||||
serde_json::json!({})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn normalizes_string_encoded_tool_input() {
|
||||
let tool = UseTool;
|
||||
@@ -902,88 +488,6 @@ mod tests {
|
||||
assert_eq!(captured, serde_json::Value::String("not json".into()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_tool_with_invalid_local_tool_id_falls_back_to_gateway() {
|
||||
let gateway_captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
|
||||
let ctx = ctx_with_dispatch_and_resources(
|
||||
NotFoundDispatch,
|
||||
gateway_resources_with_expected_call_id(
|
||||
Arc::clone(&gateway_captured),
|
||||
serde_json::json!("gateway ran"),
|
||||
Some("gateway.invalidLocal"),
|
||||
),
|
||||
);
|
||||
|
||||
let result = kigi_tool_runtime::Tool::run(
|
||||
&UseTool,
|
||||
ctx,
|
||||
UseToolInput {
|
||||
tool_name: "connector__bad/id".into(),
|
||||
tool_input: serde_json::json!({"q": "x"}),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(gateway_captured.lock().unwrap().clone().unwrap()["q"], "x");
|
||||
assert!(matches!(result, ToolOutput::MCP(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_catalog_collision_propagates_local_non_not_found_error() {
|
||||
let gateway_captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
|
||||
let ctx = ctx_with_dispatch_and_resources(
|
||||
InvalidArgumentsDispatch,
|
||||
gateway_resources(
|
||||
Arc::clone(&gateway_captured),
|
||||
serde_json::json!("gateway should not run"),
|
||||
),
|
||||
);
|
||||
|
||||
let result = kigi_tool_runtime::Tool::run(
|
||||
&UseTool,
|
||||
ctx,
|
||||
UseToolInput {
|
||||
tool_name: "server__tool".into(),
|
||||
tool_input: serde_json::json!({"local": true}),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let err = result.unwrap_err();
|
||||
assert_eq!(err.kind, kigi_tool_runtime::ToolErrorKind::InvalidArguments);
|
||||
assert!(err.detail.contains("local validation failed"));
|
||||
assert!(gateway_captured.lock().unwrap().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_catalog_collision_prefers_local_dispatch_for_server_tool() {
|
||||
let captured: SharedArgs = Arc::new(std::sync::Mutex::new(None));
|
||||
let ctx = ctx_with_dispatch_and_resources(
|
||||
CapturingDispatch {
|
||||
captured_args: Arc::clone(&captured),
|
||||
},
|
||||
gateway_resources(
|
||||
Arc::new(std::sync::Mutex::new(None)),
|
||||
serde_json::json!("gateway should not run"),
|
||||
),
|
||||
);
|
||||
|
||||
let result = kigi_tool_runtime::Tool::run(
|
||||
&UseTool,
|
||||
ctx,
|
||||
UseToolInput {
|
||||
tool_name: "server__tool".into(),
|
||||
tool_input: serde_json::json!({"local": true}),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(result.is_ok());
|
||||
let captured = captured.lock().unwrap().clone().unwrap();
|
||||
assert_eq!(captured, serde_json::json!({"local": true}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_server_tool_still_uses_local_dispatch_path() {
|
||||
let tool = UseTool;
|
||||
|
||||
@@ -105,10 +105,6 @@ pub fn canonical_input(input: &ToolInput) -> Option<serde_json::Value> {
|
||||
| ToolInput::KillTask(_)
|
||||
| ToolInput::Task(_)
|
||||
| ToolInput::WebSearch(_)
|
||||
| ToolInput::ImageGen(_)
|
||||
| ToolInput::ImageEdit(_)
|
||||
| ToolInput::ImageToVideo(_)
|
||||
| ToolInput::ReferenceToVideo(_)
|
||||
| ToolInput::WebFetch(_)
|
||||
| ToolInput::ApplyPatch(_)
|
||||
| ToolInput::HashlineEdit(_)
|
||||
|
||||
@@ -263,16 +263,6 @@ pub struct SessionContext {
|
||||
/// passed to every session. Same pattern as `fs` and `backend`.
|
||||
/// When `Some`, inserted into `Resources` so `LspTool` can use it.
|
||||
pub lsp: Option<std::sync::Arc<dyn crate::implementations::lsp::LspBackend>>,
|
||||
/// Optional image generation configuration. When `Enabled`, an `ImageGenClient`
|
||||
/// is created and injected into `Resources` so the `image_gen` tool can
|
||||
/// call the xAI Imagine API. When `Disabled` (default), the tool is not
|
||||
/// registered and image generation is unavailable.
|
||||
pub image_gen_config: crate::implementations::grok_build::image_gen::ImageGenConfig,
|
||||
/// Optional video generation configuration. When `Enabled`, a `VideoGenClient`
|
||||
/// is created and injected into `Resources` so the `video_gen` tool can
|
||||
/// call the xAI Video Generation API. When `Disabled` (default), the tool is not
|
||||
/// registered and video generation is unavailable.
|
||||
pub video_gen_config: crate::implementations::grok_build::video_gen::VideoGenConfig,
|
||||
/// Optional deploy service configuration. When enabled, the
|
||||
/// `deploy_app` tool connects to the service at call time using the shared
|
||||
/// API key provider.
|
||||
@@ -284,7 +274,7 @@ pub struct SessionContext {
|
||||
/// Prevents 401 failures when a session outlives the initial token lifetime.
|
||||
pub api_key_provider: Option<crate::types::SharedApiKeyProvider>,
|
||||
/// Optional 401-attribution callback for tool HTTP clients. When
|
||||
/// set, a 401 from `image_gen` / `video_gen` / `web_search`
|
||||
/// set, a 401 from `web_search`
|
||||
/// emits an `auth_401_attribution` event via this hook. Hosts can
|
||||
/// wire this to the same attribution sink used for inference-side
|
||||
/// 401s so tool and chat auth failures share one telemetry path.
|
||||
@@ -673,10 +663,6 @@ impl ToolRegistryBuilder {
|
||||
b.register::<grok_build::WebSearchTool>();
|
||||
b.register_with_params::<grok_build::WebFetchTool, grok_build::web_fetch::WebFetchParams>();
|
||||
b.register::<grok_build::LspTool>();
|
||||
b.register::<grok_build::ImageGenTool>();
|
||||
b.register::<grok_build::ImageEditTool>();
|
||||
b.register::<grok_build::ImageToVideoTool>();
|
||||
b.register::<grok_build::ReferenceToVideoTool>();
|
||||
b.register::<grok_build::EnterPlanModeTool>();
|
||||
b.register::<grok_build::ExitPlanModeTool>();
|
||||
b.register_with_params::<
|
||||
@@ -994,34 +980,6 @@ impl ToolRegistryBuilder {
|
||||
if let Some(lsp) = ctx.lsp {
|
||||
resources.insert(lsp);
|
||||
}
|
||||
if ctx.image_gen_config.has_credentials() {
|
||||
match crate::implementations::grok_build::image_gen::ImageGenClient::new(
|
||||
&ctx.image_gen_config,
|
||||
ctx.api_key_provider.clone(),
|
||||
) {
|
||||
Ok(client) => {
|
||||
let client = client.with_attribution_callback(ctx.attribution_callback.clone());
|
||||
resources.insert(client);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to create ImageGenClient: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
if ctx.video_gen_config.is_enabled() {
|
||||
match crate::implementations::grok_build::video_gen::VideoGenClient::new(
|
||||
&ctx.video_gen_config,
|
||||
ctx.api_key_provider.clone(),
|
||||
) {
|
||||
Ok(client) => {
|
||||
let client = client.with_attribution_callback(ctx.attribution_callback.clone());
|
||||
resources.insert(client);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to create VideoGenClient: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
if let crate::implementations::grok_build::web_fetch::WebFetchConfig::Enabled { params } =
|
||||
&ctx.web_fetch_config
|
||||
{
|
||||
@@ -1999,10 +1957,6 @@ mod tests {
|
||||
web_fetch_config:
|
||||
crate::implementations::grok_build::web_fetch::WebFetchConfig::default(),
|
||||
lsp: None,
|
||||
image_gen_config:
|
||||
crate::implementations::grok_build::image_gen::ImageGenConfig::default(),
|
||||
video_gen_config:
|
||||
crate::implementations::grok_build::video_gen::VideoGenConfig::default(),
|
||||
app_builder_deployer_config:
|
||||
crate::implementations::grok_build::deploy_app::AppBuilderDeployerConfig::default(),
|
||||
api_key_provider: None,
|
||||
@@ -2148,7 +2102,6 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn full_toolset_descriptions_render_cleanly() {
|
||||
use crate::implementations::grok_build::{
|
||||
IMAGE_GEN_TOOL_NAME, IMAGE_TO_VIDEO_TOOL_NAME, REFERENCE_TO_VIDEO_TOOL_NAME,
|
||||
SCHEDULER_CREATE_TOOL_NAME, SCHEDULER_DELETE_TOOL_NAME,
|
||||
};
|
||||
let builder = ToolRegistryBuilder::new();
|
||||
@@ -2169,9 +2122,6 @@ mod tests {
|
||||
"web_search",
|
||||
"web_fetch",
|
||||
"lsp",
|
||||
IMAGE_GEN_TOOL_NAME,
|
||||
IMAGE_TO_VIDEO_TOOL_NAME,
|
||||
REFERENCE_TO_VIDEO_TOOL_NAME,
|
||||
"monitor",
|
||||
SCHEDULER_CREATE_TOOL_NAME,
|
||||
SCHEDULER_DELETE_TOOL_NAME,
|
||||
|
||||
@@ -597,10 +597,6 @@ pub fn consumed_completion_ids(output: &ToolOutput) -> Vec<&str> {
|
||||
| ToolOutput::SchedulerDelete(_)
|
||||
| ToolOutput::SchedulerList(_)
|
||||
| ToolOutput::UpdateGoal(_)
|
||||
| ToolOutput::ImageGen(_)
|
||||
| ToolOutput::ImageToVideo(_)
|
||||
| ToolOutput::ReferenceToVideo(_)
|
||||
| ToolOutput::ImageEdit(_)
|
||||
| ToolOutput::Dynamic(_) => {}
|
||||
}
|
||||
ids
|
||||
|
||||
@@ -59,10 +59,6 @@ impl ToolKind {
|
||||
ToolKind::EnterPlan => "Enter Plan Mode",
|
||||
ToolKind::ExitPlan => "Exit Plan Mode",
|
||||
ToolKind::AskUser => "Ask User",
|
||||
ToolKind::ImageGen => "Generate Image",
|
||||
ToolKind::VideoGen => "Generate Video",
|
||||
ToolKind::ImageToVideo => "Generate Video",
|
||||
ToolKind::ReferenceToVideo => "Generate Video",
|
||||
ToolKind::DeployApp => "Deploy App",
|
||||
ToolKind::SearchTool => "Search Tools",
|
||||
ToolKind::UseTool => "Use Tool",
|
||||
@@ -100,10 +96,6 @@ impl ToolKind {
|
||||
| ToolKind::KillTaskAction
|
||||
| ToolKind::Skill
|
||||
| ToolKind::Task
|
||||
| ToolKind::ImageGen
|
||||
| ToolKind::VideoGen
|
||||
| ToolKind::ImageToVideo
|
||||
| ToolKind::ReferenceToVideo
|
||||
| ToolKind::DeployApp
|
||||
| ToolKind::SearchTool
|
||||
| ToolKind::UseTool
|
||||
|
||||
@@ -55,73 +55,6 @@ impl From<serde_json::Value> for DynamicOutput {
|
||||
Self { value }
|
||||
}
|
||||
}
|
||||
/// Typed saved path for the media tools (`image_gen` / `video_gen` /
|
||||
/// `image_edit`), so consumers read it directly instead of scraping the prose.
|
||||
/// A struct (not a bare `PathBuf`) is required: `ToolOutput` is internally
|
||||
/// tagged and only accepts map payloads.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct MediaGenOutput {
|
||||
/// Absolute path to the saved media file. Empty for [`Self::uploaded`].
|
||||
pub path: PathBuf,
|
||||
/// Basename of the saved media file (for example, `8.jpg`).
|
||||
#[serde(default)]
|
||||
pub filename: String,
|
||||
/// Session-relative media directory name (for example, `images` or `videos`).
|
||||
#[serde(default)]
|
||||
pub session_folder: String,
|
||||
/// Set when the media was uploaded to a remote presigned URL (ZDR video
|
||||
/// output) and is not available locally; omitted otherwise.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub uploaded_url: Option<String>,
|
||||
}
|
||||
impl MediaGenOutput {
|
||||
pub fn new(path: PathBuf) -> Self {
|
||||
let filename = path
|
||||
.file_name()
|
||||
.map(|name| name.to_string_lossy().into_owned())
|
||||
.unwrap_or_default();
|
||||
let session_folder = path
|
||||
.parent()
|
||||
.and_then(|parent| parent.file_name())
|
||||
.map(|name| name.to_string_lossy().into_owned())
|
||||
.unwrap_or_default();
|
||||
Self {
|
||||
path,
|
||||
filename,
|
||||
session_folder,
|
||||
uploaded_url: None,
|
||||
}
|
||||
}
|
||||
/// Media uploaded to a remote presigned URL and not available locally
|
||||
/// (ZDR video output). No local path/filename/session folder.
|
||||
pub fn uploaded(url: String) -> Self {
|
||||
Self {
|
||||
path: PathBuf::new(),
|
||||
filename: String::new(),
|
||||
session_folder: String::new(),
|
||||
uploaded_url: Some(url),
|
||||
}
|
||||
}
|
||||
/// Model-facing prose. `action` is the variant's lead-in
|
||||
/// ("Image generated" / "Video generated" / "Image edited"); the trailing
|
||||
/// guidance stops the model re-reading or narrating the result.
|
||||
pub fn prompt_text(&self, action: &str) -> String {
|
||||
if let Some(url) = &self.uploaded_url {
|
||||
return format!(
|
||||
"{action} and uploaded to {url}. The file is not available locally — reference it by this URL. Do not read or re-display it, and do not describe how it appears to the user."
|
||||
);
|
||||
}
|
||||
let path = self.path.to_string_lossy().to_string();
|
||||
let message = format!(
|
||||
"{action} and saved to {path}. Do not read or re-display it, and do not describe how it appears to the user."
|
||||
);
|
||||
serde_json::json!(
|
||||
{ "path" : path, "filename" : & self.filename, "session_folder" : & self
|
||||
.session_folder, "message" : message, }
|
||||
)
|
||||
.to_string()
|
||||
}
|
||||
}
|
||||
use crate::implementations::grok_build::todo::{TodoItem, TodoState};
|
||||
use crate::implementations::skills::skill::SkillOutput;
|
||||
use crate::util::truncate::{DEFAULT_SOFT_WRAP_WIDTH, soft_wrap_lines};
|
||||
@@ -653,14 +586,6 @@ pub enum ToolOutput {
|
||||
/// (e.g., memory_search, memory_get). The string is the pre-formatted
|
||||
/// prompt text — no additional rendering is needed.
|
||||
Text(TextOutput),
|
||||
#[from(skip)]
|
||||
ImageGen(MediaGenOutput),
|
||||
#[from(skip)]
|
||||
ImageToVideo(MediaGenOutput),
|
||||
#[from(skip)]
|
||||
ReferenceToVideo(MediaGenOutput),
|
||||
#[from(skip)]
|
||||
ImageEdit(MediaGenOutput),
|
||||
}
|
||||
impl ToolOutput {
|
||||
/// Whether this output is a logical tool failure, for `tool.execution`'s
|
||||
@@ -977,10 +902,6 @@ impl ToolOutput {
|
||||
ToolOutput::UpdateGoal(o) => o.summary.clone(),
|
||||
ToolOutput::Dynamic(v) => serde_json::to_string_pretty(&v.value).unwrap_or_default(),
|
||||
ToolOutput::Text(text) => text.text.clone(),
|
||||
ToolOutput::ImageGen(m) => m.prompt_text("Image generated"),
|
||||
ToolOutput::ImageToVideo(m) => m.prompt_text("Video generated"),
|
||||
ToolOutput::ReferenceToVideo(m) => m.prompt_text("Video generated"),
|
||||
ToolOutput::ImageEdit(m) => m.prompt_text("Image edited"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1320,91 +1241,6 @@ mod tests {
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn media_gen_output() {
|
||||
let cases = [
|
||||
(
|
||||
ToolOutput::ImageGen(MediaGenOutput::new("/tmp/images/1.jpg".into())),
|
||||
"ImageGen",
|
||||
"/tmp/images/1.jpg",
|
||||
"1.jpg",
|
||||
"images",
|
||||
"Image generated and saved to /tmp/images/1.jpg. Do not read or re-display it, and do not describe how it appears to the user.",
|
||||
),
|
||||
(
|
||||
ToolOutput::ImageToVideo(MediaGenOutput::new("/tmp/videos/2.mp4".into())),
|
||||
"ImageToVideo",
|
||||
"/tmp/videos/2.mp4",
|
||||
"2.mp4",
|
||||
"videos",
|
||||
"Video generated and saved to /tmp/videos/2.mp4. Do not read or re-display it, and do not describe how it appears to the user.",
|
||||
),
|
||||
(
|
||||
ToolOutput::ReferenceToVideo(MediaGenOutput::new("/tmp/videos/3.mp4".into())),
|
||||
"ReferenceToVideo",
|
||||
"/tmp/videos/3.mp4",
|
||||
"3.mp4",
|
||||
"videos",
|
||||
"Video generated and saved to /tmp/videos/3.mp4. Do not read or re-display it, and do not describe how it appears to the user.",
|
||||
),
|
||||
(
|
||||
ToolOutput::ImageEdit(MediaGenOutput::new("/tmp/images/2.jpg".into())),
|
||||
"ImageEdit",
|
||||
"/tmp/images/2.jpg",
|
||||
"2.jpg",
|
||||
"images",
|
||||
"Image edited and saved to /tmp/images/2.jpg. Do not read or re-display it, and do not describe how it appears to the user.",
|
||||
),
|
||||
];
|
||||
for (output, ty, path, filename, session_folder, message) in cases {
|
||||
let prompt_json: serde_json::Value =
|
||||
serde_json::from_str(&output.to_prompt_format()).unwrap();
|
||||
assert_eq!(prompt_json["path"], path);
|
||||
assert_eq!(prompt_json["filename"], filename);
|
||||
assert_eq!(prompt_json["session_folder"], session_folder);
|
||||
assert_eq!(prompt_json["message"], message);
|
||||
let json = to_json(output);
|
||||
assert_eq!(json["type"], ty);
|
||||
assert_eq!(json["path"], path);
|
||||
assert_eq!(json["filename"], filename);
|
||||
assert_eq!(json["session_folder"], session_folder);
|
||||
let (ToolOutput::ImageGen(m)
|
||||
| ToolOutput::ImageToVideo(m)
|
||||
| ToolOutput::ReferenceToVideo(m)
|
||||
| ToolOutput::ImageEdit(m)) = serde_json::from_value(json).unwrap()
|
||||
else {
|
||||
panic!("unexpected variant");
|
||||
};
|
||||
assert_eq!(m.path, PathBuf::from(path));
|
||||
assert_eq!(m.filename, filename);
|
||||
assert_eq!(m.session_folder, session_folder);
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn media_gen_output_uploaded() {
|
||||
let url = "https://files.example.com/team/video-abc.mp4";
|
||||
let output = ToolOutput::ImageToVideo(MediaGenOutput::uploaded(url.to_string()));
|
||||
let prompt = output.to_prompt_format();
|
||||
assert!(prompt.contains(url), "prompt must include the upload URL");
|
||||
assert!(
|
||||
prompt.contains("not available locally"),
|
||||
"prompt must tell the model the file is remote-only"
|
||||
);
|
||||
assert!(
|
||||
prompt.contains("Do not read or re-display"),
|
||||
"prompt must include re-display guard"
|
||||
);
|
||||
let json = to_json(output);
|
||||
assert_eq!(json["uploaded_url"], url);
|
||||
assert!(
|
||||
json.get("path").is_some(),
|
||||
"path field must be present (empty for uploaded)"
|
||||
);
|
||||
let ToolOutput::ImageToVideo(m) = serde_json::from_value(json).unwrap() else {
|
||||
panic!("unexpected variant");
|
||||
};
|
||||
assert_eq!(m, MediaGenOutput::uploaded(url.to_string()));
|
||||
}
|
||||
#[test]
|
||||
fn read_file_not_found_json() {
|
||||
let json =
|
||||
to_json(ReadFileOutput::FileNotFound("Error: /tmp/x does not exist.".into()).into());
|
||||
|
||||
@@ -544,42 +544,6 @@ pub fn display_cwd_or_cwd(cwd: &std::path::Path, display_cwd: Option<&std::path:
|
||||
/// through the outer `ToolBridge` (which would deadlock).
|
||||
#[derive(Clone)]
|
||||
pub struct InnerDispatch(pub std::sync::Arc<dyn kigi_tool_runtime::ToolDispatch>);
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ManagedGatewayToolSource {
|
||||
pub connector_id: String,
|
||||
pub connector_name: String,
|
||||
pub tool_id: String,
|
||||
pub tool_name: String,
|
||||
pub call_id: String,
|
||||
}
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ManagedGatewayToolCatalog(pub HashMap<String, ManagedGatewayToolSource>);
|
||||
impl ManagedGatewayToolCatalog {
|
||||
pub fn get(&self, name: &str) -> Option<&ManagedGatewayToolSource> {
|
||||
self.0.get(name)
|
||||
}
|
||||
}
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ManagedGatewayToolCallResponse {
|
||||
pub result: serde_json::Value,
|
||||
pub connectors_needing_reauth: Vec<String>,
|
||||
}
|
||||
#[async_trait::async_trait]
|
||||
pub trait ManagedGatewayToolCaller: Send + Sync {
|
||||
async fn call_tool(
|
||||
&self,
|
||||
call_id: &str,
|
||||
arguments: serde_json::Value,
|
||||
caller: &str,
|
||||
) -> Result<ManagedGatewayToolCallResponse, kigi_tool_runtime::ToolError>;
|
||||
}
|
||||
#[derive(Clone)]
|
||||
pub struct ManagedGatewayToolClient(pub Arc<dyn ManagedGatewayToolCaller>);
|
||||
impl std::fmt::Debug for ManagedGatewayToolClient {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("ManagedGatewayToolClient").finish()
|
||||
}
|
||||
}
|
||||
/// Whether streaming output is enabled for this invocation.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct StreamEnabled(pub bool);
|
||||
@@ -753,33 +717,6 @@ impl AvailableSkills {
|
||||
/// Session folder for logs and output files.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionFolder(pub PathBuf);
|
||||
/// Per-turn registry mapping each attached image's `[Image #N]` display
|
||||
/// number to a reference `image_edit` can resolve.
|
||||
///
|
||||
/// The model sees attachments inline (as pixels) and only the `[Image #N]`
|
||||
/// token in text — never a path — so this lets `image_edit` resolve that
|
||||
/// token instead of fabricating a filesystem path it can't know.
|
||||
///
|
||||
/// Keyed by display **number**, not list position: numbers are not
|
||||
/// renumbered when a chip is removed mid-compose (`#1` and `#3` survive
|
||||
/// after `#2`) and images may be dropped during normalization, so the two
|
||||
/// diverge. Each reference is a bare filesystem path (the durable
|
||||
/// `session_image_path`) or a `data:<mime>;base64,<data>` URL fallback.
|
||||
///
|
||||
/// Replaced wholesale each turn (empty when there are no attachments) so a
|
||||
/// stale registry never resolves to a prior turn's image. Ephemeral — not
|
||||
/// persisted, not serde-registered.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct AttachedImages(pub Vec<(usize, String)>);
|
||||
impl AttachedImages {
|
||||
/// Resolve an `[Image #N]` display number to its reference string.
|
||||
pub fn reference_for(&self, display_number: usize) -> Option<&str> {
|
||||
self.0
|
||||
.iter()
|
||||
.find(|(n, _)| *n == display_number)
|
||||
.map(|(_, reference)| reference.as_str())
|
||||
}
|
||||
}
|
||||
/// Notification handle for streaming tool output.
|
||||
#[derive(Clone)]
|
||||
pub struct NotificationHandle(pub ToolNotificationHandle);
|
||||
|
||||
@@ -91,10 +91,6 @@ pub enum ToolKind {
|
||||
EnterPlan,
|
||||
ExitPlan,
|
||||
AskUser,
|
||||
ImageGen,
|
||||
VideoGen,
|
||||
ImageToVideo,
|
||||
ReferenceToVideo,
|
||||
DeployApp,
|
||||
SearchTool,
|
||||
UseTool,
|
||||
|
||||
@@ -19,14 +19,11 @@ use crate::implementations::grok_build::ask_user_question::AskUserQuestionInput;
|
||||
use crate::implementations::grok_build::enter_plan_mode::EnterPlanModeInput;
|
||||
use crate::implementations::grok_build::exit_plan_mode::ExitPlanModeInput;
|
||||
use crate::implementations::grok_build::grep::GrepSearchInput;
|
||||
use crate::implementations::grok_build::image_edit::ImageEditInput;
|
||||
use crate::implementations::grok_build::image_gen::ImageGenInput;
|
||||
use crate::implementations::grok_build::list_dir::ListDirInput;
|
||||
use crate::implementations::grok_build::read_file::ReadFileInput;
|
||||
use crate::implementations::grok_build::search_replace::SearchReplaceInput;
|
||||
use crate::implementations::grok_build::todo::TodoWriteInput;
|
||||
use crate::implementations::grok_build::update_goal::UpdateGoalInput;
|
||||
use crate::implementations::grok_build::video_gen::{ImageToVideoInput, ReferenceToVideoInput};
|
||||
use crate::implementations::grok_build::web_fetch::WebFetchInput;
|
||||
use crate::implementations::grok_build::web_search::WebSearchInput;
|
||||
use crate::implementations::lsp::LspToolInput;
|
||||
@@ -71,10 +68,6 @@ pub enum ToolInput {
|
||||
KillTask(KillTaskToolInput),
|
||||
Task(TaskToolInput),
|
||||
WebSearch(WebSearchInput),
|
||||
ImageGen(ImageGenInput),
|
||||
ImageEdit(ImageEditInput),
|
||||
ImageToVideo(ImageToVideoInput),
|
||||
ReferenceToVideo(ReferenceToVideoInput),
|
||||
WebFetch(WebFetchInput),
|
||||
Write(WriteInput),
|
||||
ApplyPatch(ApplyPatchInput),
|
||||
|
||||
Reference in New Issue
Block a user