@samitouri / QOSami_Agents_Zero-2 / commits / 4f06aa0a

Fix MCP multimodal content handling

Preserve MCP image, audio, and resource tool results instead of collapsing non-text responses into an empty textual result. Images and image resources now flow into raw history as data URL attachments, while audio and non-image binary resources are saved as artifacts with normalized paths. Extract shared media artifact helpers for base64 validation, image data URLs, decoded-size checks, artifact saving, MIME normalization, and safe filenames. Reuse the shared helpers from MCP, browser connector, and computer-use artifact paths, and add focused regression coverage.

Alessandro committed May 26, 2026 at 15:31 UTC 4f06aa0a8eae842f295cc8304519f08f7fb7812b

6 files changed +859 -81

helpers/mcp_handler.py

+259 -50

@@ -21,6 +21,7 @@ from contextlib import AsyncExitStack

21 from shutil import which

22 from datetime import timedelta

23 import json

24 +import uuid

25 from helpers import errors

26 from helpers import settings

27 from helpers.log import LogItem

@@ -39,11 +40,21 @@ from anyio.streams.memory import (

40 )

42 from pydantic import BaseModel, Field, Discriminator, Tag, PrivateAttr

42 -from helpers import dirty_json

43 +from helpers import dirty_json, media_artifacts

44 from helpers.print_style import PrintStyle

45 from helpers.tool import Tool, Response

48 +MCP_MEDIA_TOKENS_ESTIMATE = 1500

49 +MAX_MCP_RESOURCE_TEXT_CHARS = 12_000

50 +

51 +

52 +def _mcp_get(item: Any, key: str, default: Any = None) -> Any:

53 + if isinstance(item, dict):

54 + return item.get(key, default)

55 + return getattr(item, key, default)

56 +

57 +

58 def normalize_name(name: str) -> str:

59 # Lowercase and strip whitespace

60 name = name.strip().lower()

@@ -102,7 +113,6 @@ class MCPTool(Tool):

113 """MCP Tool wrapper"""

114

115 def get_log_object(self) -> LogItem:

105 - import uuid

116 return self.agent.context.log.log(

117 type="mcp",

118 heading=f"icon://extension {self.agent.agent_name}: Using MCP tool '{self.name}'",

@@ -111,17 +121,235 @@ class MCPTool(Tool):

121 id=str(uuid.uuid4()),

122 )

123

124 + def _context_id(self) -> str:

125 + return str(getattr(getattr(self.agent, "context", None), "id", "") or "").strip()

126 +

127 + def _raw_tool_response(self, response: Response) -> str:

128 + raw_tool_response = response.message.strip() if response.message else ""

129 + if not raw_tool_response:

130 + PrintStyle(font_color="red").print(

131 + f"Warning: Tool '{self.name}' returned an empty message."

132 + )

133 + raw_tool_response = "[Tool returned no textual content]"

134 + return raw_tool_response

135 +

136 + def _coerce_media_token_estimate(self, value: object) -> int:

137 + try:

138 + estimate = int(value or 0)

139 + except (TypeError, ValueError):

140 + estimate = 0

141 + return estimate if estimate > 0 else MCP_MEDIA_TOKENS_ESTIMATE

142 +

143 + def _format_image_content(

144 + self,

145 + *,

146 + encoded: str,

147 + mime_type: str,

148 + label: str,

149 + ) -> tuple[str, dict[str, Any] | None]:

150 + try:

151 + image = media_artifacts.image_data_url_from_base64(

152 + encoded,

153 + mime_type=mime_type,

154 + )

155 + except media_artifacts.EmptyBase64Data:

156 + return f"MCP returned an empty {label} attachment.", None

157 + except media_artifacts.InvalidBase64Data:

158 + return f"MCP returned a {label} attachment that could not be decoded.", None

159 +

160 + return (

161 + f"MCP returned {label} attachment ({image.mime}, {image.size} bytes).",

162 + {

163 + "type": "image_url",

164 + "image_url": {"url": image.url},

165 + },

166 + )

167 +

168 + def _materialize_binary_content(

169 + self,

170 + *,

171 + encoded: str,

172 + mime_type: str,

173 + label: str,

174 + index: int,

175 + preferred_name: str = "",

176 + ) -> str:

177 + try:

178 + safe_mime = media_artifacts.normalize_mime(mime_type)

179 + artifact = media_artifacts.save_base64_artifact(

180 + encoded,

181 + mime_type=safe_mime,

182 + directory_parts=self._artifact_directory_parts(),

183 + preferred_name=preferred_name,

184 + default_filename=self._default_artifact_filename(

185 + label=label,

186 + index=index,

187 + mime_type=safe_mime,

188 + ),

189 + )

190 + except media_artifacts.EmptyBase64Data:

191 + return f"MCP returned an empty {label} attachment."

192 + except media_artifacts.InvalidBase64Data:

193 + return f"MCP returned a {label} attachment that could not be decoded."

194 +

195 + return f"Saved MCP {label} attachment ({artifact.mime}, {artifact.size} bytes) to {artifact.path}."

196 +

197 + def _artifact_directory_parts(self) -> tuple[str, ...]:

198 + context_id = normalize_name(self._context_id() or "shared") or "shared"

199 + tool_name = normalize_name(self.name or "mcp_tool") or "mcp_tool"

200 + return ("tmp", "mcp", context_id, tool_name)

201 +

202 + def _default_artifact_filename(self, *, label: str, index: int, mime_type: str) -> str:

203 + tool_name = normalize_name(self.name or "mcp_tool") or "mcp_tool"

204 + ext = media_artifacts.guess_extension(mime_type, ".bin")

205 + return f"{tool_name}_{label}_{index}{ext}"

206 +

207 + def _format_resource_text(self, text: str, uri: str = "") -> str:

208 + body = str(text or "").strip()

209 + if not body:

210 + return ""

211 + if len(body) > MAX_MCP_RESOURCE_TEXT_CHARS:

212 + body = body[:MAX_MCP_RESOURCE_TEXT_CHARS].rstrip() + "\n...[truncated]"

213 + if uri:

214 + return f"Resource {uri}:\n{body}"

215 + return body

216 +

217 + def _content_item_dump(self, item: Any) -> dict[str, Any]:

218 + if isinstance(item, dict):

219 + return dict(item)

220 + model_dump = getattr(item, "model_dump", None)

221 + if callable(model_dump):

222 + dumped = model_dump(mode="python")

223 + if isinstance(dumped, dict):

224 + return dumped

225 + item_vars = getattr(item, "__dict__", None)

226 + if isinstance(item_vars, dict):

227 + return dict(item_vars)

228 + return {}

229 +

230 + def _summarize_unknown_item(self, item: Any, item_type: str) -> str:

231 + dumped = self._content_item_dump(item)

232 + if dumped:

233 + dumped.pop("data", None)

234 + resource = dumped.get("resource")

235 + if isinstance(resource, dict):

236 + resource.pop("blob", None)

237 + summary = json.dumps(dumped, ensure_ascii=False)

238 + if len(summary) > 600:

239 + summary = summary[:600] + "...[truncated]"

240 + return f"MCP returned unsupported content item type '{item_type}': {summary}"

241 + return f"MCP returned unsupported content item type '{item_type}'."

242 +

243 + def _format_tool_result(

244 + self, response: CallToolResult

245 + ) -> tuple[str, dict[str, Any] | None]:

246 + text_parts: list[str] = []

247 + notes: list[str] = []

248 + raw_images: list[dict[str, Any]] = []

249 + content_items = list(getattr(response, "content", []) or [])

250 +

251 + for index, item in enumerate(content_items, start=1):

252 + item_type = str(_mcp_get(item, "type", "") or "").strip().lower()

253 +

254 + if item_type == "text":

255 + text = str(_mcp_get(item, "text", "") or "").strip()

256 + if text:

257 + text_parts.append(text)

258 + continue

259 +

260 + if item_type == "image":

261 + note, raw_content = self._format_image_content(

262 + encoded=str(_mcp_get(item, "data", "") or ""),

263 + mime_type=str(_mcp_get(item, "mimeType", "") or "image/png"),

264 + label="image",

265 + )

266 + notes.append(note)

267 + if raw_content:

268 + raw_images.append(raw_content)

269 + continue

270 +

271 + if item_type == "audio":

272 + note = self._materialize_binary_content(

273 + encoded=str(_mcp_get(item, "data", "") or ""),

274 + mime_type=str(_mcp_get(item, "mimeType", "") or "audio/wav"),

275 + label="audio",

276 + index=index,

277 + )

278 + notes.append(note)

279 + continue

280 +

281 + if item_type == "resource":

282 + resource = _mcp_get(item, "resource", None)

283 + uri = str(_mcp_get(resource, "uri", "") or "").strip()

284 + text = _mcp_get(resource, "text", None)

285 + if isinstance(text, str) and text.strip():

286 + text_parts.append(self._format_resource_text(text, uri))

287 + continue

288 +

289 + blob = str(_mcp_get(resource, "blob", "") or "").strip()

290 + if blob:

291 + mime_type = str(

292 + _mcp_get(resource, "mimeType", "") or "application/octet-stream"

293 + ).strip().lower()

294 + if mime_type.startswith("image/"):

295 + note, raw_content = self._format_image_content(

296 + encoded=blob,

297 + mime_type=mime_type,

298 + label="resource image",

299 + )

300 + else:

301 + note = self._materialize_binary_content(

302 + encoded=blob,

303 + mime_type=mime_type,

304 + label="resource",

305 + index=index,

306 + preferred_name=uri,

307 + )

308 + raw_content = None

309 + notes.append(note)

310 + if raw_content:

311 + raw_images.append(raw_content)

312 + continue

313 +

314 + if uri:

315 + mime_type = str(_mcp_get(resource, "mimeType", "") or "").strip()

316 + details = f" ({mime_type})" if mime_type else ""

317 + notes.append(f"MCP returned a resource reference: {uri}{details}.")

318 + continue

319 +

320 + notes.append("MCP returned a resource item without text or binary data.")

321 + continue

322 +

323 + if item_type:

324 + notes.append(self._summarize_unknown_item(item, item_type))

325 + continue

326 +

327 + notes.append(self._summarize_unknown_item(item, "unknown"))

328 +

329 + message = "\n\n".join(part for part in [*text_parts, *notes] if part.strip())

330 + if not message and content_items:

331 + message = "MCP tool returned content that could not be rendered as text."

332 +

333 + additional = None

334 + if raw_images:

335 + additional = {

336 + "raw_content": raw_images,

337 + "preview": f"<MCP image attachments: {len(raw_images)}>",

338 + "_tokens": MCP_MEDIA_TOKENS_ESTIMATE * len(raw_images),

339 + }

340 +

341 + return message, additional

342 +

343 async def execute(self, **kwargs: Any):

344 error = ""

345 + additional: dict[str, Any] | None = None

346 try:

347 response: CallToolResult = await MCPConfig.get_instance().call_tool(

348 self.name, kwargs

349 )

120 - message = "\n\n".join(

121 - [item.text for item in response.content if item.type == "text"]

122 - )

350 + message, additional = self._format_tool_result(response)

351 if response.isError:

124 - error = message

352 + error = message or "MCP tool returned an error without textual content."

353 except Exception as e:

354 error = f"MCP Tool Exception: {str(e)}"

355 message = f"ERROR: {str(e)}"

@@ -139,7 +367,7 @@ class MCPTool(Tool):

367 content=f"{self.name}: {error}",

368 )

369

142 - return Response(message=message, break_loop=False)

370 + return Response(message=message, break_loop=False, additional=additional)

371

372 async def before_execution(self, **kwargs: Any):

373 (

@@ -159,48 +387,29 @@ class MCPTool(Tool):

387 PrintStyle().print()

388

389 async def after_execution(self, response: Response, **kwargs: Any):

162 - raw_tool_response = response.message.strip() if response.message else ""

163 - if not raw_tool_response:

164 - PrintStyle(font_color="red").print(

165 - f"Warning: Tool '{self.name}' returned an empty message."

390 + final_text_for_agent = self._raw_tool_response(response)

391 + additional = dict(response.additional or {})

392 + raw_content = additional.pop("raw_content", None)

393 + preview = str(additional.pop("preview", "") or "").strip()

394 + token_estimate = self._coerce_media_token_estimate(additional.pop("_tokens", 0))

395 +

396 + self.agent.hist_add_tool_result(

397 + self.name,

398 + final_text_for_agent,

399 + id=self.log.id if self.log else "",

400 + **additional,

401 + )

402 + if raw_content:

403 + from helpers import history

404 +

405 + self.agent.hist_add_message(

406 + False,

407 + content=history.RawMessage(

408 + raw_content=raw_content,

409 + preview=preview or final_text_for_agent,

410 + ),

411 + tokens=token_estimate,

412 )

167 - # Even if empty, we might still want to provide context for the agent

168 - raw_tool_response = "[Tool returned no textual content]"

169 -

170 - # Prepare user message context

171 - # user_message_text = (

172 - # "No specific user message context available for this exact step."

173 - # )

174 - # if (

175 - # self.agent

176 - # and self.agent.last_user_message

177 - # and self.agent.last_user_message.content

178 - # ):

179 - # content = self.agent.last_user_message.content

180 - # if isinstance(content, dict):

181 - # # Attempt to get a 'message' field, otherwise stringify the dict

182 - # user_message_text = str(content.get(

183 - # "message", json.dumps(content, indent=2)

184 - # ))

185 - # elif isinstance(content, str):

186 - # user_message_text = content

187 - # else:

188 - # # Fallback for any other types (e.g. list, if that were possible for content)

189 - # user_message_text = str(content)

190 -

191 - # # Ensure user_message_text is a string before length check and slicing

192 - # user_message_text = str(user_message_text)

193 -

194 - # # Truncate user message context if it's too long to avoid overwhelming the prompt

195 - # max_user_context_len = 500 # characters

196 - # if len(user_message_text) > max_user_context_len:

197 - # user_message_text = (

198 - # user_message_text[:max_user_context_len] + "... (truncated)"

199 - # )

200 -

201 - final_text_for_agent = raw_tool_response

202 -

203 - self.agent.hist_add_tool_result(self.name, final_text_for_agent, id=self.log.id if self.log else "")

413 (

414 PrintStyle(

415 font_color="#1B4F72", background_color="white", padding=True, bold=True

@@ -210,8 +419,8 @@ class MCPTool(Tool):

419 )

420 # Print only the raw response to console for brevity, agent gets the full context.

421 PrintStyle(font_color="#85C1E9").print(

213 - raw_tool_response

214 - if raw_tool_response

422 + final_text_for_agent

423 + if final_text_for_agent

424 else "[No direct textual output from tool]"

425 )

426 if self.log:

helpers/media_artifacts.py new

+186

@@ -0,0 +1,186 @@

1 +from __future__ import annotations

2 +

3 +import base64

4 +import binascii

5 +import mimetypes

6 +import uuid

7 +from dataclasses import dataclass

8 +from pathlib import Path

9 +from urllib.parse import urlparse

10 +

11 +from helpers import files

12 +

13 +

14 +DEFAULT_MAX_ARTIFACT_SIZE_BYTES = 25 * 1024 * 1024

15 +

16 +

17 +class MediaArtifactError(ValueError):

18 + pass

19 +

20 +

21 +class EmptyBase64Data(MediaArtifactError):

22 + pass

23 +

24 +

25 +class InvalidBase64Data(MediaArtifactError):

26 + pass

27 +

28 +

29 +class ArtifactTooLarge(MediaArtifactError):

30 + def __init__(self, size: int, limit: int):

31 + super().__init__(f"artifact is too large ({size} bytes, limit {limit} bytes)")

32 + self.size = size

33 + self.limit = limit

34 +

35 +

36 +@dataclass(frozen=True)

37 +class Base64Payload:

38 + data: str

39 + payload: bytes

40 + size: int

41 +

42 +

43 +@dataclass(frozen=True)

44 +class ImageDataUrl:

45 + url: str

46 + mime: str

47 + size: int

48 +

49 +

50 +@dataclass(frozen=True)

51 +class SavedArtifact:

52 + path: str

53 + mime: str

54 + size: int

55 +

56 +

57 +def compact_base64(data: str) -> str:

58 + return "".join(char for char in str(data or "") if not char.isspace())

59 +

60 +

61 +def estimated_base64_decoded_size(data: str) -> int:

62 + compact_length = len(compact_base64(data))

63 + return (compact_length * 3) // 4

64 +

65 +

66 +def decode_base64_payload(

67 + data: str,

68 + *,

69 + max_bytes: int | None = None,

70 +) -> Base64Payload:

71 + compact = compact_base64(data)

72 + if not compact:

73 + raise EmptyBase64Data("base64 data is empty")

74 +

75 + if max_bytes is not None:

76 + estimated_size = estimated_base64_decoded_size(compact)

77 + if estimated_size > max_bytes:

78 + raise ArtifactTooLarge(estimated_size, max_bytes)

79 +

80 + try:

81 + payload = base64.b64decode(compact, validate=True)

82 + except (binascii.Error, ValueError) as exc:

83 + raise InvalidBase64Data("base64 data could not be decoded") from exc

84 +

85 + if max_bytes is not None and len(payload) > max_bytes:

86 + raise ArtifactTooLarge(len(payload), max_bytes)

87 +

88 + return Base64Payload(data=compact, payload=payload, size=len(payload))

89 +

90 +

91 +def normalize_mime(

92 + mime_type: str,

93 + *,

94 + default: str = "application/octet-stream",

95 + required_prefix: str = "",

96 +) -> str:

97 + value = str(mime_type or "").strip().lower()

98 + if not value:

99 + return default

100 + if required_prefix and not value.startswith(required_prefix):

101 + return default

102 + return value

103 +

104 +

105 +def guess_extension(mime_type: str, fallback: str = ".bin") -> str:

106 + ext = mimetypes.guess_extension(str(mime_type or "").strip().lower()) or fallback

107 + return ".jpg" if ext == ".jpe" else ext

108 +

109 +

110 +def filename_from_uri(uri: str) -> str:

111 + value = str(uri or "").strip()

112 + if not value:

113 + return ""

114 + parsed = urlparse(value)

115 + return Path(parsed.path or value).name

116 +

117 +

118 +def safe_filename(

119 + value: str,

120 + *,

121 + default: str = "artifact.bin",

122 + default_extension: str = ".bin",

123 +) -> str:

124 + source = str(value or "").strip() or default

125 + cleaned = "".join(

126 + char if char.isalnum() or char in {"-", "_", "."} else "_"

127 + for char in source

128 + )

129 + cleaned = cleaned.strip("._") or default

130 + if "." not in cleaned:

131 + ext = default_extension if default_extension.startswith(".") else f".{default_extension}"

132 + cleaned += ext

133 + return cleaned

134 +

135 +

136 +def image_data_url_from_base64(

137 + data: str,

138 + *,

139 + mime_type: str = "image/png",

140 + max_bytes: int | None = None,

141 +) -> ImageDataUrl:

142 + payload = decode_base64_payload(data, max_bytes=max_bytes)

143 + safe_mime = normalize_mime(

144 + mime_type,

145 + default="image/png",

146 + required_prefix="image/",

147 + )

148 + return ImageDataUrl(

149 + url=f"data:{safe_mime};base64,{payload.data}",

150 + mime=safe_mime,

151 + size=payload.size,

152 + )

153 +

154 +

155 +def save_base64_artifact(

156 + data: str,

157 + *,

158 + mime_type: str = "application/octet-stream",

159 + directory_parts: tuple[str, ...],

160 + preferred_name: str = "",

161 + default_filename: str = "artifact.bin",

162 + max_bytes: int | None = None,

163 +) -> SavedArtifact:

164 + payload = decode_base64_payload(data, max_bytes=max_bytes)

165 + safe_mime = normalize_mime(mime_type)

166 + preferred_filename = filename_from_uri(preferred_name) or default_filename

167 + default_extension = guess_extension(safe_mime, Path(default_filename).suffix or ".bin")

168 + filename = safe_filename(

169 + preferred_filename,

170 + default=default_filename,

171 + default_extension=default_extension,

172 + )

173 + filename_path = Path(filename)

174 + stem = filename_path.stem or Path(default_filename).stem or "artifact"

175 + suffix = filename_path.suffix or default_extension

176 +

177 + artifact_dir = Path(files.get_abs_path(*directory_parts))

178 + artifact_dir.mkdir(parents=True, exist_ok=True)

179 + path = artifact_dir / f"{stem}_{uuid.uuid4().hex[:8]}{suffix}"

180 + path.write_bytes(payload.payload)

181 +

182 + return SavedArtifact(

183 + path=files.normalize_a0_path(str(path)),

184 + mime=safe_mime,

185 + size=payload.size,

186 + )

plugins/_a0_connector/tools/computer_use_remote.py

+7 -16

@@ -6,7 +6,7 @@ from pathlib import Path

6 import uuid

7 from typing import Any

9 -from helpers import history

9 +from helpers import history, media_artifacts

10 from helpers.print_style import PrintStyle

11 from helpers.tool import Response, Tool

12 from helpers.ws import NAMESPACE

@@ -750,7 +750,7 @@ class ComputerUseRemote(Tool):

750 if isinstance(artifact, dict) and str(artifact.get("encoding", "")).strip().lower() == "base64":

751 encoded = str(artifact.get("data") or "")

752 if encoded:

753 - estimated_size = _estimated_base64_decoded_size(encoded)

753 + estimated_size = media_artifacts.estimated_base64_decoded_size(encoded)

754 if estimated_size > MAX_CAPTURE_ARTIFACT_SIZE_BYTES:

755 raise RuntimeError(

756 "Computer-use capture artifact is too large to attach safely "

@@ -759,7 +759,11 @@ class ComputerUseRemote(Tool):

759 mime = str(artifact.get("mime") or "image/png").strip()

760 if not mime.startswith("image/"):

761 mime = "image/png"

762 - filename = _safe_filename(str(artifact.get("filename") or "computer-use-capture.png"))

762 + filename = media_artifacts.safe_filename(

763 + str(artifact.get("filename") or "computer-use-capture.png"),

764 + default=f"computer-use-{uuid.uuid4().hex}.png",

765 + default_extension=".png",

766 + )

767 return f"data:{mime};base64,{encoded}", Path(filename).stem

768

769 if path_error is not None:

@@ -845,19 +849,6 @@ class ComputerUseRemote(Tool):

849 return str(value or "").strip().lower() in {"1", "true", "yes", "on"}

850

851

848 -def _safe_filename(value: str) -> str:

849 - cleaned = "".join(char if char.isalnum() or char in {"-", "_", "."} else "_" for char in value)

850 - cleaned = cleaned.strip("._") or f"computer-use-{uuid.uuid4().hex}.png"

851 - if "." not in cleaned:

852 - cleaned += ".png"

853 - return cleaned

854 -

855 -

856 -def _estimated_base64_decoded_size(data: str) -> int:

857 - compact_length = sum(1 for char in data if not char.isspace())

858 - return (compact_length * 3) // 4

859 -

860 -

852 def _sanitize_tool_text(value: str) -> str:

853 try:

854 from helpers.strings import sanitize_string

plugins/_browser/helpers/connector_runtime.py

+7 -15

@@ -9,7 +9,7 @@ from pathlib import Path

9 from typing import Any

10 from urllib.parse import urlparse

12 -from helpers import ephemeral_images

12 +from helpers import ephemeral_images, media_artifacts

14 try:

15 from helpers.ws import NAMESPACE

@@ -425,13 +425,17 @@ class ConnectorBrowserRuntime:

425 data = str(artifact.get("data") or "")

426 if not data:

427 return result

428 - estimated_size = _estimated_base64_decoded_size(data)

428 + estimated_size = media_artifacts.estimated_base64_decoded_size(data)

429 if estimated_size > MAX_ARTIFACT_SIZE_BYTES:

430 raise RuntimeError(

431 "Host browser artifact is too large to attach safely "

432 f"({estimated_size} bytes, limit {MAX_ARTIFACT_SIZE_BYTES} bytes)."

433 )

434 - filename = _safe_filename(str(artifact.get("filename") or "host-browser.jpg"))

434 + filename = media_artifacts.safe_filename(

435 + str(artifact.get("filename") or "host-browser.jpg"),

436 + default=f"host-browser-{uuid.uuid4().hex}.jpg",

437 + default_extension=".jpg",

438 + )

439 try:

440 ref = ephemeral_images.put_image(

441 context_id=self.context_id,

@@ -558,15 +562,3 @@ def _api_base_is_local(api_base: str) -> bool:

562 hostname = (parsed.hostname or "").strip().lower()

563 return hostname in _LOCAL_HOSTS

564

561 -

562 -def _safe_filename(value: str) -> str:

563 - cleaned = "".join(char if char.isalnum() or char in {"-", "_", "."} else "_" for char in value)

564 - cleaned = cleaned.strip("._") or f"host-browser-{uuid.uuid4().hex}.jpg"

565 - if "." not in cleaned:

566 - cleaned += ".jpg"

567 - return cleaned

568 -

569 -

570 -def _estimated_base64_decoded_size(data: str) -> int:

571 - compact_length = sum(1 for char in data if not char.isspace())

572 - return (compact_length * 3) // 4

tests/test_mcp_handler_multimodal.py new

+334

@@ -0,0 +1,334 @@

1 +from __future__ import annotations

2 +

3 +import asyncio

4 +import base64

5 +import importlib

6 +import sys

7 +from dataclasses import dataclass

8 +from pathlib import Path

9 +from types import ModuleType, SimpleNamespace

10 +

11 +import pytest

12 +

13 +

14 +PROJECT_ROOT = Path(__file__).resolve().parents[1]

15 +if str(PROJECT_ROOT) not in sys.path:

16 + sys.path.insert(0, str(PROJECT_ROOT))

17 +

18 +

19 +@dataclass

20 +class _StubResponse:

21 + message: str

22 + break_loop: bool

23 + additional: dict | None = None

24 +

25 +

26 +class _StubTool:

27 + def __init__(

28 + self,

29 + agent=None,

30 + name="",

31 + method=None,

32 + args=None,

33 + message="",

34 + loop_data=None,

35 + **kwargs,

36 + ):

37 + self.agent = agent

38 + self.name = name

39 + self.method = method

40 + self.args = args or {}

41 + self.message = message

42 + self.loop_data = loop_data

43 + self.log = None

44 +

45 + def nice_key(self, key: str) -> str:

46 + return key

47 +

48 +

49 +class _FakeContent(SimpleNamespace):

50 + pass

51 +

52 +

53 +class _FakeCallToolResult(SimpleNamespace):

54 + pass

55 +

56 +

57 +@pytest.fixture

58 +def mcp_handler_module(monkeypatch, tmp_path):

59 + monkeypatch.delitem(sys.modules, "helpers.mcp_handler", raising=False)

60 +

61 + agent_module = ModuleType("agent")

62 + agent_module.AgentContext = type("AgentContext", (), {})

63 + agent_module.Agent = type("Agent", (), {})

64 + agent_module.LoopData = type("LoopData", (), {})

65 + monkeypatch.setitem(sys.modules, "agent", agent_module)

66 +

67 + tool_module = ModuleType("helpers.tool")

68 + tool_module.Response = _StubResponse

69 + tool_module.Tool = _StubTool

70 + monkeypatch.setitem(sys.modules, "helpers.tool", tool_module)

71 +

72 + settings_module = ModuleType("helpers.settings")

73 + monkeypatch.setitem(sys.modules, "helpers.settings", settings_module)

74 +

75 + history_module = ModuleType("helpers.history")

76 + history_module.RawMessage = lambda **kwargs: dict(kwargs)

77 + monkeypatch.setitem(sys.modules, "helpers.history", history_module)

78 +

79 + mcp_module = ModuleType("mcp")

80 + mcp_module.ClientSession = type("ClientSession", (), {})

81 + mcp_module.StdioServerParameters = type("StdioServerParameters", (), {})

82 + monkeypatch.setitem(sys.modules, "mcp", mcp_module)

83 +

84 + mcp_client_stdio = ModuleType("mcp.client.stdio")

85 + mcp_client_stdio.stdio_client = lambda *args, **kwargs: None

86 + monkeypatch.setitem(sys.modules, "mcp.client.stdio", mcp_client_stdio)

87 +

88 + mcp_client_sse = ModuleType("mcp.client.sse")

89 + mcp_client_sse.sse_client = lambda *args, **kwargs: None

90 + monkeypatch.setitem(sys.modules, "mcp.client.sse", mcp_client_sse)

91 +

92 + mcp_client_streamable_http = ModuleType("mcp.client.streamable_http")

93 + mcp_client_streamable_http.streamablehttp_client = lambda *args, **kwargs: None

94 + monkeypatch.setitem(

95 + sys.modules,

96 + "mcp.client.streamable_http",

97 + mcp_client_streamable_http,

98 + )

99 +

100 + mcp_shared_message = ModuleType("mcp.shared.message")

101 + mcp_shared_message.SessionMessage = type("SessionMessage", (), {})

102 + monkeypatch.setitem(sys.modules, "mcp.shared.message", mcp_shared_message)

103 +

104 + mcp_types = ModuleType("mcp.types")

105 + mcp_types.CallToolResult = _FakeCallToolResult

106 + mcp_types.ListToolsResult = type("ListToolsResult", (), {})

107 + monkeypatch.setitem(sys.modules, "mcp.types", mcp_types)

108 +

109 + module = importlib.import_module("helpers.mcp_handler")

110 +

111 + class _SilentPrintStyle:

112 + def __init__(self, *args, **kwargs):

113 + pass

114 +

115 + def print(self, *args, **kwargs):

116 + return self

117 +

118 + def stream(self, *args, **kwargs):

119 + return self

120 +

121 + def _fake_get_abs_path(*parts):

122 + return str(tmp_path.joinpath(*parts))

123 +

124 + def _fake_normalize_a0_path(path: str) -> str:

125 + path_obj = Path(path)

126 + try:

127 + rel = path_obj.relative_to(tmp_path)

128 + except ValueError:

129 + return str(path_obj)

130 + return "/a0/" + str(rel).replace("\\", "/")

131 +

132 + monkeypatch.setattr(module, "PrintStyle", _SilentPrintStyle)

133 + monkeypatch.setattr(module.media_artifacts.files, "get_abs_path", _fake_get_abs_path)

134 + monkeypatch.setattr(module.media_artifacts.files, "normalize_a0_path", _fake_normalize_a0_path)

135 + return module, tmp_path

136 +

137 +

138 +def _agent_recorder(context_id: str = "ctx-mcp"):

139 + tool_results: list[tuple[tuple, dict]] = []

140 + messages: list[tuple[tuple, dict]] = []

141 + updates: list[dict] = []

142 + warnings: list[dict] = []

143 + agent = SimpleNamespace(

144 + agent_name="Agent Zero",

145 + context=SimpleNamespace(

146 + id=context_id,

147 + log=SimpleNamespace(log=lambda **kwargs: warnings.append(kwargs)),

148 + ),

149 + hist_add_tool_result=lambda *args, **kwargs: tool_results.append((args, kwargs)),

150 + hist_add_message=lambda *args, **kwargs: messages.append((args, kwargs)),

151 + )

152 + log = SimpleNamespace(id="mcp-log", update=lambda **kwargs: updates.append(kwargs))

153 + return agent, log, tool_results, messages, updates, warnings

154 +

155 +

156 +def test_mcp_image_content_becomes_history_image_attachment(mcp_handler_module, monkeypatch):

157 + module, _tmp_path = mcp_handler_module

158 + agent, log, tool_results, messages, updates, warnings = _agent_recorder()

159 + image_b64 = base64.b64encode(b"image-bytes").decode("ascii")

160 + result = _FakeCallToolResult(

161 + content=[_FakeContent(type="image", data=image_b64, mimeType="image/webp")],

162 + isError=False,

163 + )

164 +

165 + class _FakeConfig:

166 + async def call_tool(self, name, kwargs):

167 + return result

168 +

169 + monkeypatch.setattr(module.MCPConfig, "get_instance", lambda: _FakeConfig())

170 +

171 + tool = module.MCPTool(

172 + agent=agent,

173 + name="venice_image",

174 + method=None,

175 + args={},

176 + message="",

177 + loop_data=None,

178 + )

179 + tool.log = log

180 +

181 + response = asyncio.run(tool.execute())

182 +

183 + assert "[Tool returned no textual content]" not in response.message

184 + assert response.message == "MCP returned image attachment (image/webp, 11 bytes)."

185 + assert response.additional is not None

186 + data_url = response.additional["raw_content"][0]["image_url"]["url"]

187 + assert data_url == f"data:image/webp;base64,{image_b64}"

188 +

189 + asyncio.run(tool.after_execution(response))

190 +

191 + assert tool_results[0][0] == ("venice_image", response.message)

192 + raw_message = messages[0][1]["content"]

193 + assert raw_message["raw_content"][0]["image_url"]["url"] == data_url

194 + assert messages[0][1]["tokens"] == module.MCP_MEDIA_TOKENS_ESTIMATE

195 + assert updates[-1]["content"] == response.message

196 + assert warnings == []

197 +

198 +

199 +def test_mcp_audio_content_is_saved_instead_of_discarded(mcp_handler_module, monkeypatch):

200 + module, tmp_path = mcp_handler_module

201 + agent, log, tool_results, messages, updates, warnings = _agent_recorder()

202 + audio_b64 = base64.b64encode(b"audio-bytes").decode("ascii")

203 + result = _FakeCallToolResult(

204 + content=[_FakeContent(type="audio", data=audio_b64, mimeType="audio/mpeg")],

205 + isError=False,

206 + )

207 +

208 + class _FakeConfig:

209 + async def call_tool(self, name, kwargs):

210 + return result

211 +

212 + monkeypatch.setattr(module.MCPConfig, "get_instance", lambda: _FakeConfig())

213 +

214 + tool = module.MCPTool(

215 + agent=agent,

216 + name="venice_audio",

217 + method=None,

218 + args={},

219 + message="",

220 + loop_data=None,

221 + )

222 + tool.log = log

223 +

224 + response = asyncio.run(tool.execute())

225 +

226 + assert response.additional is None

227 + assert "[Tool returned no textual content]" not in response.message

228 + assert "Saved MCP audio attachment (audio/mpeg, 11 bytes) to /a0/tmp/mcp/ctx_mcp/venice_audio/" in response.message

229 + saved_path = response.message.split(" to ", 1)[1].rstrip(".")

230 + assert (tmp_path / saved_path.removeprefix("/a0/")).exists()

231 +

232 + asyncio.run(tool.after_execution(response))

233 +

234 + assert tool_results[0][0] == ("venice_audio", response.message)

235 + assert messages == []

236 + assert updates[-1]["content"] == response.message

237 + assert warnings == []

238 +

239 +

240 +def test_mcp_image_resource_blob_becomes_history_image_attachment(mcp_handler_module, monkeypatch):

241 + module, _tmp_path = mcp_handler_module

242 + agent, log, tool_results, messages, updates, warnings = _agent_recorder()

243 + image_b64 = base64.b64encode(b"resource-image").decode("ascii")

244 + result = _FakeCallToolResult(

245 + content=[

246 + _FakeContent(

247 + type="resource",

248 + resource=_FakeContent(

249 + uri="memory://venice/image.webp",

250 + mimeType="image/webp",

251 + blob=image_b64,

252 + ),

253 + )

254 + ],

255 + isError=False,

256 + )

257 +

258 + class _FakeConfig:

259 + async def call_tool(self, name, kwargs):

260 + return result

261 +

262 + monkeypatch.setattr(module.MCPConfig, "get_instance", lambda: _FakeConfig())

263 +

264 + tool = module.MCPTool(

265 + agent=agent,

266 + name="venice_resource_image",

267 + method=None,

268 + args={},

269 + message="",

270 + loop_data=None,

271 + )

272 + tool.log = log

273 +

274 + response = asyncio.run(tool.execute())

275 +

276 + assert response.message == "MCP returned resource image attachment (image/webp, 14 bytes)."

277 + assert response.additional is not None

278 + data_url = response.additional["raw_content"][0]["image_url"]["url"]

279 + assert data_url == f"data:image/webp;base64,{image_b64}"

280 +

281 + asyncio.run(tool.after_execution(response))

282 +

283 + assert tool_results[0][0] == ("venice_resource_image", response.message)

284 + raw_message = messages[0][1]["content"]

285 + assert raw_message["raw_content"][0]["image_url"]["url"] == data_url

286 + assert updates[-1]["content"] == response.message

287 + assert warnings == []

288 +

289 +

290 +def test_mcp_resource_text_is_preserved(mcp_handler_module, monkeypatch):

291 + module, _tmp_path = mcp_handler_module

292 + agent, log, tool_results, messages, updates, warnings = _agent_recorder()

293 + result = _FakeCallToolResult(

294 + content=[

295 + _FakeContent(

296 + type="resource",

297 + resource=_FakeContent(

298 + uri="memory://venice/caption.txt",

299 + mimeType="text/plain",

300 + text="Generated caption text",

301 + ),

302 + )

303 + ],

304 + isError=False,

305 + )

306 +

307 + class _FakeConfig:

308 + async def call_tool(self, name, kwargs):

309 + return result

310 +

311 + monkeypatch.setattr(module.MCPConfig, "get_instance", lambda: _FakeConfig())

312 +

313 + tool = module.MCPTool(

314 + agent=agent,

315 + name="venice_resource",

316 + method=None,

317 + args={},

318 + message="",

319 + loop_data=None,

320 + )

321 + tool.log = log

322 +

323 + response = asyncio.run(tool.execute())

324 +

325 + assert response.additional is None

326 + assert "Resource memory://venice/caption.txt:" in response.message

327 + assert "Generated caption text" in response.message

328 +

329 + asyncio.run(tool.after_execution(response))

330 +

331 + assert tool_results[0][0] == ("venice_resource", response.message)

332 + assert messages == []

333 + assert updates[-1]["content"] == response.message

334 + assert warnings == []

tests/test_media_artifacts.py new

+66

@@ -0,0 +1,66 @@

1 +from __future__ import annotations

2 +

3 +import base64

4 +import sys

5 +from pathlib import Path

6 +

7 +import pytest

8 +

9 +

10 +PROJECT_ROOT = Path(__file__).resolve().parents[1]

11 +if str(PROJECT_ROOT) not in sys.path:

12 + sys.path.insert(0, str(PROJECT_ROOT))

13 +

14 +from helpers import media_artifacts

15 +

16 +

17 +def test_image_data_url_from_base64_compacts_and_normalizes_mime():

18 + encoded = " \n" + base64.b64encode(b"image-bytes").decode("ascii") + "\n"

19 +

20 + image = media_artifacts.image_data_url_from_base64(

21 + encoded,

22 + mime_type="IMAGE/WEBP",

23 + )

24 +

25 + assert image.mime == "image/webp"

26 + assert image.size == 11

27 + assert image.url == "data:image/webp;base64,aW1hZ2UtYnl0ZXM="

28 +

29 +

30 +def test_decode_base64_payload_rejects_empty_invalid_and_oversized_data():

31 + with pytest.raises(media_artifacts.EmptyBase64Data):

32 + media_artifacts.decode_base64_payload(" \n ")

33 +

34 + with pytest.raises(media_artifacts.InvalidBase64Data):

35 + media_artifacts.decode_base64_payload("not base64")

36 +

37 + with pytest.raises(media_artifacts.ArtifactTooLarge) as exc_info:

38 + media_artifacts.decode_base64_payload("ZmFrZQ==", max_bytes=2)

39 +

40 + assert exc_info.value.size == media_artifacts.estimated_base64_decoded_size("ZmFrZQ==")

41 + assert exc_info.value.limit == 2

42 +

43 +

44 +def test_save_base64_artifact_uses_uri_filename_and_normalized_a0_path(monkeypatch, tmp_path):

45 + def fake_get_abs_path(*parts):

46 + return str(tmp_path.joinpath(*parts))

47 +

48 + def fake_normalize_a0_path(path: str):

49 + return "/a0/" + str(Path(path).relative_to(tmp_path)).replace("\\", "/")

50 +

51 + monkeypatch.setattr(media_artifacts.files, "get_abs_path", fake_get_abs_path)

52 + monkeypatch.setattr(media_artifacts.files, "normalize_a0_path", fake_normalize_a0_path)

53 +

54 + artifact = media_artifacts.save_base64_artifact(

55 + base64.b64encode(b"audio-bytes").decode("ascii"),

56 + mime_type="audio/mpeg",

57 + directory_parts=("tmp", "media-test"),

58 + preferred_name="memory://venice/generated track.mp3",

59 + default_filename="fallback.bin",

60 + )

61 +

62 + assert artifact.mime == "audio/mpeg"

63 + assert artifact.size == 11

64 + assert artifact.path.startswith("/a0/tmp/media-test/generated_track_")

65 + assert artifact.path.endswith(".mp3")

66 + assert (tmp_path / artifact.path.removeprefix("/a0/")).read_bytes() == b"audio-bytes"