Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,21 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[BaseEvent]:
run_id = input_data.run_id or str(uuid.uuid4())
result_key = (thread_id, run_id)

# Validate and convert the prompt before touching the thread's live
# SessionWorker. Unsupported content is an input error, not evidence
# that the existing Claude CLI session is broken.
try:
prompt, _ = process_messages(input_data)
except Exception as e:
logger.error(f"Invalid input for thread={thread_id}: {e}")
yield RunErrorEvent(
type=EventType.RUN_ERROR,
thread_id=thread_id,
run_id=run_id,
message=str(e),
)
return

# ── Run-admission serialization (Fix 1) ──
# Acquire the per-thread RUN lock at admission — BEFORE worker.query() /
# RUN_STARTED — and hold it across the WHOLE run, releasing in the
Expand Down Expand Up @@ -345,7 +360,6 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[BaseEvent]:
worker = entry["worker"]
logger.debug(f"Reusing worker for thread={thread_id}")

prompt, _ = process_messages(input_data)
message_stream = worker.query(prompt, session_id=thread_id)

# Log parent_run_id if provided (for branching/time travel tracking)
Expand Down Expand Up @@ -1158,4 +1172,3 @@ def flush_pending_msg():
type=EventType.MESSAGES_SNAPSHOT,
messages=all_messages,
)

Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import asyncio
import logging
from contextlib import suppress
from collections.abc import AsyncIterable
from typing import Any, AsyncIterator, Optional

logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -162,7 +163,11 @@ async def _graceful_disconnect(client: Any) -> None:
except Exception as exc:
logger.debug(f"[SessionWorker] Graceful disconnect error (ignored): {exc}")

async def query(self, prompt: str, session_id: str = "default") -> AsyncIterator[Any]:
async def query(
self,
prompt: str | AsyncIterable[dict[str, Any]],
session_id: str = "default",
) -> AsyncIterator[Any]:
"""Send prompt to the worker and yield SDK Message objects."""
output_queue: asyncio.Queue = asyncio.Queue()
# Register the output queue in the in-flight set BEFORE enqueuing the
Expand Down
193 changes: 179 additions & 14 deletions integrations/claude-agent-sdk/python/ag_ui_claude_sdk/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,23 @@

import json
import logging
from collections.abc import AsyncIterable, AsyncIterator
from typing import Any, Dict, List, Optional, Tuple
from ag_ui.core import RunAgentInput, AssistantMessage, ToolCall, FunctionCall, ToolMessage
from ag_ui.core import (
RunAgentInput,
AssistantMessage,
ToolCall,
FunctionCall,
ToolMessage,
TextInputContent,
ImageInputContent,
AudioInputContent,
VideoInputContent,
DocumentInputContent,
BinaryInputContent,
InputContentDataSource,
InputContentUrlSource,
)

from .config import STATE_MANAGEMENT_TOOL_NAME, STATE_MANAGEMENT_TOOL_FULL_NAME

Expand Down Expand Up @@ -89,7 +104,150 @@ def strip_mcp_prefix(tool_name: str) -> str:
return tool_name


def process_messages(input_data: RunAgentInput) -> Tuple[str, bool]:
SUPPORTED_IMAGE_MEDIA_TYPES = {
"image/jpeg",
"image/png",
"image/gif",
"image/webp",
}


def _normalized_media_type(value: Optional[str]) -> Optional[str]:
if not isinstance(value, str):
return None
media_type = value.split(";", 1)[0].strip().lower()
return media_type or None


def _require_non_empty_string(value: Any, field: str) -> str:
if not isinstance(value, str) or not value:
raise ValueError(f"{field} must be a non-empty string")
return value


def _require_remote_url(value: Any, field: str) -> str:
from urllib.parse import urlsplit

url = _require_non_empty_string(value, field)
parsed = urlsplit(url)
if parsed.scheme not in {"http", "https"} or not parsed.netloc:
raise ValueError(f"{field} must be a valid http or https URL")
return url


def _image_block(source: Any, field: str) -> Dict[str, Any]:
media_type = _normalized_media_type(getattr(source, "mime_type", None))
if isinstance(source, InputContentDataSource):
if media_type not in SUPPORTED_IMAGE_MEDIA_TYPES:
raise ValueError(
f"{field}.mime_type must be image/jpeg, image/png, image/gif, or image/webp"
)
return {
"type": "image",
"source": {
"type": "base64",
"media_type": media_type,
"data": _require_non_empty_string(source.value, f"{field}.value"),
},
}
if isinstance(source, InputContentUrlSource):
if media_type is not None and media_type not in SUPPORTED_IMAGE_MEDIA_TYPES:
raise ValueError(f"{field}.mime_type is not a supported image type")
return {
"type": "image",
"source": {
"type": "url",
"url": _require_remote_url(source.value, f"{field}.value"),
},
}
raise ValueError(f"{field} must be a data or URL source")


def _document_block(source: Any, field: str) -> Dict[str, Any]:
media_type = _normalized_media_type(getattr(source, "mime_type", None))
if isinstance(source, InputContentDataSource):
if media_type != "application/pdf":
raise ValueError(f"{field}.mime_type must be application/pdf")
return {
"type": "document",
"source": {
"type": "base64",
"media_type": "application/pdf",
"data": _require_non_empty_string(source.value, f"{field}.value"),
},
}
if isinstance(source, InputContentUrlSource):
if media_type is not None and media_type != "application/pdf":
raise ValueError(f"{field}.mime_type must be application/pdf when provided")
return {
"type": "document",
"source": {
"type": "url",
"url": _require_remote_url(source.value, f"{field}.value"),
},
}
raise ValueError(f"{field} must be a data or URL source")


def _legacy_binary_block(block: BinaryInputContent, index: int) -> Dict[str, Any]:
media_type = _normalized_media_type(block.mime_type)
if block.data:
source: Any = InputContentDataSource(
value=block.data,
mime_type=media_type or "",
)
elif block.url:
source = InputContentUrlSource(
value=block.url,
mime_type=media_type,
)
else:
raise ValueError(
f"content[{index}] uses an opaque file id, which the Claude Agent SDK adapter cannot resolve"
)

if media_type in SUPPORTED_IMAGE_MEDIA_TYPES:
return _image_block(source, f"content[{index}]")
if media_type == "application/pdf":
return _document_block(source, f"content[{index}]")
raise ValueError(f"content[{index}].mime_type is not supported")


def _convert_content_block(block: Any, index: int) -> Optional[Dict[str, Any]]:
if isinstance(block, TextInputContent):
if not block.text.strip():
return None
return {
"type": "text",
"text": block.text,
}
if isinstance(block, ImageInputContent):
return _image_block(block.source, f"content[{index}].source")
if isinstance(block, DocumentInputContent):
return _document_block(block.source, f"content[{index}].source")
if isinstance(block, BinaryInputContent):
return _legacy_binary_block(block, index)
if isinstance(block, (AudioInputContent, VideoInputContent)):
raise ValueError(f"content[{index}] type {block.type} is not supported")
raise ValueError(f"content[{index}] has an unsupported type")


async def _structured_user_message(
content: List[Dict[str, Any]],
session_id: str,
) -> AsyncIterator[Dict[str, Any]]:
yield {
"type": "user",
"message": {"role": "user", "content": content},
"parent_tool_use_id": None,
"session_id": session_id,
}


ClaudePrompt = str | AsyncIterable[Dict[str, Any]]


def process_messages(input_data: RunAgentInput) -> Tuple[ClaudePrompt, bool]:
"""
Process and validate all messages from RunAgentInput.

Expand All @@ -100,7 +258,9 @@ def process_messages(input_data: RunAgentInput) -> Tuple[str, bool]:
input_data: RunAgentInput with messages array

Returns:
Tuple of (user_message: str, has_pending_tool_result: bool)
Tuple of (user_message, has_pending_tool_result). ``user_message`` is
a string for plain text or a one-message async iterable for structured
content.
"""
messages = input_data.messages or []

Expand Down Expand Up @@ -133,7 +293,8 @@ def process_messages(input_data: RunAgentInput) -> Tuple[str, bool]:

# Extract content from the LAST message (any role - user, tool, or assistant)
# Claude SDK manages conversation history via session_id, we just need the latest input
user_message = ""
user_message: ClaudePrompt = ""
has_user_content = False
if messages:
last_msg = messages[-1]

Expand All @@ -148,17 +309,21 @@ def process_messages(input_data: RunAgentInput) -> Tuple[str, bool]:
# Handle different content formats
if isinstance(content, str):
user_message = content
has_user_content = bool(content)
elif isinstance(content, list):
# Content blocks format - extract text from first text block
for block in content:
if hasattr(block, 'text'):
user_message = block.text
break
elif isinstance(block, dict) and 'text' in block:
user_message = block['text']
break

if not user_message:
blocks = []
for index, block in enumerate(content):
converted = _convert_content_block(block, index)
if converted is not None:
blocks.append(converted)
if blocks:
user_message = _structured_user_message(
blocks,
input_data.thread_id or "default",
)
has_user_content = True

if not has_user_content:
logger.warning(f"No user message found in {len(messages)} messages")

return user_message, has_pending_tool_result
Expand Down
Loading
Loading