feat(xai): surface hosted X Search activity (#5050)

This commit is contained in:
chengyongru
2026-07-23 13:42:09 +08:00
committed by GitHub
parent f3099286ea
commit 9cf2fb19c2
14 changed files with 431 additions and 8 deletions
+15
View File
@@ -90,6 +90,14 @@ class AgentHook:
async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None:
pass
async def on_provider_tool_event(
self,
context: AgentHookContext,
event: dict[str, Any],
) -> None:
"""Observe a provider-hosted tool lifecycle event."""
pass
async def before_execute_tools(self, context: AgentHookContext) -> None:
pass
@@ -192,6 +200,13 @@ class CompositeHook(AgentHook):
async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None:
await self._for_each_hook_safe("on_stream_end", context, resuming=resuming)
async def on_provider_tool_event(
self,
context: AgentHookContext,
event: dict[str, Any],
) -> None:
await self._for_each_hook_safe("on_provider_tool_event", context, event)
async def before_execute_tools(self, context: AgentHookContext) -> None:
await self._for_each_hook_safe("before_execute_tools", context)
+57
View File
@@ -9,6 +9,7 @@ from typing import Any, Awaitable, Callable
from loguru import logger
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.providers.base import ToolCallRequest
from nanobot.utils.helpers import IncrementalThinkExtractor, strip_think
from nanobot.utils.progress_events import (
build_tool_event_finish_payloads,
@@ -97,6 +98,61 @@ class AgentProgressHook(AgentHook):
self._session_key,
)
async def on_provider_tool_event(
self,
context: AgentHookContext,
event: dict[str, Any],
) -> None:
if not self._on_progress:
return
phase = event.get("phase")
name = event.get("name")
call_id = event.get("call_id")
if (
phase not in {"start", "end", "error"}
or not isinstance(name, str)
or not name
or not call_id
):
return
arguments = event.get("arguments")
if not isinstance(arguments, dict):
arguments = {}
payload = {
"version": 1,
"phase": phase,
"call_id": str(call_id),
"name": name,
"arguments": arguments,
"result": event.get("result") if phase == "end" else None,
"error": event.get("error") if phase == "error" else None,
"files": [],
"embeds": [],
}
if phase == "start":
await self.emit_reasoning_end()
tool_call = ToolCallRequest(id=str(call_id), name=name, arguments=arguments)
tool_hint = self._strip_think(self._tool_hint([tool_call])) or name
await invoke_on_progress(
self._on_progress,
tool_hint,
tool_hint=True,
tool_events=[payload],
)
logger.info(
"Provider-hosted tool call: {}({})",
name,
json.dumps(arguments, ensure_ascii=False)[:200],
)
return
if on_progress_accepts_tool_events(self._on_progress):
await invoke_on_progress(
self._on_progress,
"",
tool_hint=False,
tool_events=[payload],
)
async def before_execute_tools(self, context: AgentHookContext) -> None:
if self._on_progress:
if not self._on_stream and not context.streamed_content:
@@ -114,6 +170,7 @@ class AgentProgressHook(AgentHook):
for tc in context.tool_calls:
args_str = json.dumps(tc.arguments, ensure_ascii=False)
logger.info("Tool call: {}({})", tc.name, args_str[:200])
async def emit_reasoning(self, reasoning_content: str | None) -> None:
"""Publish a reasoning chunk; channel plugins decide whether to render."""
if (
+27
View File
@@ -722,6 +722,20 @@ class AgentRunner:
)
progress_state: dict[str, bool] | None = None
active_hosted_tools: dict[str, dict[str, Any]] = {}
async def _provider_tool_event(event: dict[str, Any]) -> None:
if event.get("kind") != "hosted_tool":
return
await hook.on_provider_tool_event(context, event)
call_id = event.get("call_id")
if not call_id:
return
call_id = str(call_id)
if event.get("phase") == "start":
active_hosted_tools[call_id] = dict(event)
elif event.get("phase") in {"end", "error"}:
active_hosted_tools.pop(call_id, None)
if wants_streaming:
thinking_buf = ""
@@ -750,6 +764,7 @@ class AgentRunner:
**kwargs,
on_content_delta=_stream,
on_thinking_delta=_thinking,
on_tool_call_delta=_provider_tool_event,
on_stream_recover=_stream_recover,
)
elif wants_progress_streaming:
@@ -780,6 +795,7 @@ class AgentRunner:
coro = spec.runtime.provider.chat_stream_with_retry(
**kwargs,
on_content_delta=_stream_progress,
on_tool_call_delta=_provider_tool_event,
)
else:
coro = spec.runtime.provider.chat_with_retry(**kwargs)
@@ -813,6 +829,17 @@ class AgentRunner:
finish_reason="error",
error_kind="timeout",
)
# chat_stream_with_retry may recover internally, so only fail unfinished
# hosted calls after the provider returns its final error response.
if response.finish_reason == "error":
for event in list(active_hosted_tools.values()):
await _provider_tool_event({
**event,
"phase": "error",
"result": None,
"error": response.content
or "Model request failed before the provider-hosted tool completed.",
})
if progress_state and progress_state.get("reasoning_open"):
await hook.emit_reasoning_end()
dropped, all_dropped, original_finish_reason = (
@@ -117,6 +117,7 @@ async def consume_sse_with_reasoning(
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
on_reasoning_delta: Callable[[str], Awaitable[None]] | None = None,
on_response_event: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]:
"""Consume a Responses API SSE stream, including visible reasoning summaries."""
content = ""
@@ -129,6 +130,8 @@ async def consume_sse_with_reasoning(
streamed_reasoning = False
async for event in iter_sse(response):
if on_response_event:
await on_response_event(event)
event_type = event.get("type")
if event_type == "response.output_item.added":
item = event.get("item") or {}
+60
View File
@@ -401,6 +401,11 @@ async def _request_xai(
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]:
async def _on_response_event(event: dict[str, Any]) -> None:
hosted_event = _xai_hosted_tool_event(event)
if hosted_event is not None and on_tool_call_delta is not None:
await on_tool_call_delta(hosted_event)
client_kwargs: dict[str, Any] = {"timeout": resolve_stream_idle_timeout_s()}
if proxy:
client_kwargs.update(proxy=proxy, trust_env=False)
@@ -415,9 +420,64 @@ async def _request_xai(
on_content_delta=on_content_delta,
on_tool_call_delta=on_tool_call_delta,
on_reasoning_delta=on_thinking_delta,
on_response_event=_on_response_event if on_tool_call_delta else None,
)
def _xai_hosted_tool_event(event: dict[str, Any]) -> dict[str, Any] | None:
event_type = event.get("type")
if event_type == "response.custom_tool_call_input.done":
call_id = event.get("item_id") or event.get("call_id") or event.get("id")
if not call_id:
return None
return {
"kind": "hosted_tool",
"phase": "start",
"call_id": str(call_id),
"name": "x_search",
"arguments": _xai_hosted_tool_arguments(
event.get("input", event.get("arguments"))
),
"result": None,
}
if event_type != "response.output_item.done":
return None
item = event.get("item")
if not isinstance(item, dict) or item.get("type") != "custom_tool_call":
return None
tool_name = item.get("name")
if not isinstance(tool_name, str) or not tool_name.startswith("x_"):
return None
call_id = item.get("id") or item.get("call_id") or event.get("item_id")
if not call_id:
return None
return {
"kind": "hosted_tool",
"phase": "end",
"call_id": str(call_id),
"name": "x_search",
"arguments": _xai_hosted_tool_arguments(
item.get("input", item.get("arguments"))
),
# Keep the useful search subtype, but do not persist large hosted results
# in WebUI activity messages. The model answer already carries citations.
"result": {"name": tool_name},
}
def _xai_hosted_tool_arguments(value: Any) -> dict[str, Any]:
if isinstance(value, dict):
return dict(value)
if not isinstance(value, str) or not value.strip():
return {}
try:
parsed = json.loads(value)
except (TypeError, ValueError):
return {}
return parsed if isinstance(parsed, dict) else {}
def _build_xai_http_error(
status_code: int,
headers: httpx.Headers,
+1
View File
@@ -16,6 +16,7 @@ _TOOL_FORMATS: dict[str, tuple[list[str], str, bool, bool]] = {
"exec": (["command"], "$ {}", False, True),
"list_exec_sessions": ([], "exec sessions", False, False),
"web_search": (["query"], 'search "{}"', False, False),
"x_search": (["query"], 'search X "{}"', False, False),
"web_fetch": (["url"], "fetch {}", True, False),
"list_dir": (["path"], "ls {}", True, False),
}
+147
View File
@@ -6,7 +6,9 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from agent.runner_helpers import make_run_spec
from nanobot.agent.hook import CompositeHook
from nanobot.agent.hooks import FileEditActivityHook
from nanobot.agent.progress_hook import AgentProgressHook
from nanobot.agent.runner import AgentRunner
from nanobot.agent.tools.filesystem import EditFileTool, WriteFileTool
from nanobot.config.schema import AgentDefaults
@@ -83,6 +85,151 @@ async def test_runner_streams_provider_progress_deltas_by_default():
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_routes_hosted_tool_events_to_structured_progress():
provider = MagicMock()
provider.supports_progress_deltas = True
async def chat_stream_with_retry(*, on_content_delta, on_tool_call_delta, **kwargs):
await on_tool_call_delta({
"call_id": "local-call",
"name": "read_file",
"arguments_delta": "",
})
await on_tool_call_delta({
"kind": "hosted_tool",
"phase": "start",
"call_id": "x-search-1",
"name": "x_search",
"arguments": {"query": "nanobot oauth"},
"result": None,
})
await on_tool_call_delta({
"kind": "hosted_tool",
"phase": "end",
"call_id": "x-search-1",
"name": "x_search",
"arguments": {"query": "nanobot oauth"},
"result": {"name": "x_semantic_search"},
})
await on_content_delta("done")
return LLMResponse(content="done", tool_calls=[], usage={})
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
tools = MagicMock()
tools.get_definitions.return_value = []
progress_events: list[dict] = []
progress_text: list[str] = []
async def progress_cb(content, *, tool_events=None, **kwargs):
progress_text.append(content)
if tool_events:
progress_events.extend(tool_events)
hook = CompositeHook([AgentProgressHook(on_progress=progress_cb)])
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "search X"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
progress_callback=progress_cb,
hook=hook,
))
assert result.final_content == "done"
assert result.tools_used == []
assert result.tool_events == []
assert progress_events == [
{
"version": 1,
"phase": "start",
"call_id": "x-search-1",
"name": "x_search",
"arguments": {"query": "nanobot oauth"},
"result": None,
"error": None,
"files": [],
"embeds": [],
},
{
"version": 1,
"phase": "end",
"call_id": "x-search-1",
"name": "x_search",
"arguments": {"query": "nanobot oauth"},
"result": {"name": "x_semantic_search"},
"error": None,
"files": [],
"embeds": [],
},
]
assert progress_text == ['search X "nanobot oauth"', "", "done"]
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_fails_pending_hosted_tool_when_model_request_fails():
provider = MagicMock()
provider.supports_progress_deltas = True
async def chat_stream_with_retry(*, on_tool_call_delta, **kwargs):
await on_tool_call_delta({
"kind": "hosted_tool",
"phase": "start",
"call_id": "x-search-failed",
"name": "x_search",
"arguments": {"query": "nanobot oauth"},
"result": None,
})
return LLMResponse(
content="hosted search backend failed",
finish_reason="error",
)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
tools = MagicMock()
tools.get_definitions.return_value = []
progress_events: list[dict] = []
async def progress_cb(content, *, tool_events=None, **kwargs):
if tool_events:
progress_events.extend(tool_events)
hook = CompositeHook([AgentProgressHook(on_progress=progress_cb)])
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "search X"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
progress_callback=progress_cb,
hook=hook,
))
assert result.stop_reason == "error"
assert [(event["phase"], event["call_id"]) for event in progress_events] == [
("start", "x-search-failed"),
("error", "x-search-failed"),
]
assert progress_events[-1] == {
"version": 1,
"phase": "error",
"call_id": "x-search-failed",
"name": "x_search",
"arguments": {"query": "nanobot oauth"},
"result": None,
"error": "hosted search backend failed",
"files": [],
"embeds": [],
}
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_emits_write_file_diff_from_tool_execution_snapshots(tmp_path):
provider = MagicMock()
+68 -1
View File
@@ -354,6 +354,73 @@ async def test_raw_response_request_streams_text_usage_and_inline_citations(monk
assert captured["json"]["tools"] == [{"type": "x_search"}]
@pytest.mark.asyncio
async def test_raw_response_request_streams_hosted_x_search_lifecycle(monkeypatch) -> None:
original_client = httpx.AsyncClient
events = [
{
"type": "response.custom_tool_call_input.done",
"item_id": "x-search-1",
"input": '{"query":"nanobot oauth"}',
},
{
"type": "response.output_item.done",
"item": {
"type": "custom_tool_call",
"id": "x-search-1",
"name": "x_semantic_search",
"input": '{"query":"nanobot oauth"}',
"output": [{"text": "large hosted result must not enter activity events"}],
},
},
{
"type": "response.completed",
"response": {"status": "completed", "usage": {}},
},
]
content = "".join(f"data: {json.dumps(event)}\n\n" for event in events)
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, content=content, request=request)
def fake_client(**kwargs) -> httpx.AsyncClient:
return original_client(
transport=httpx.MockTransport(handler),
timeout=kwargs["timeout"],
)
monkeypatch.setattr("nanobot.providers.xai_grok_provider.httpx.AsyncClient", fake_client)
tool_events: list[dict[str, Any]] = []
result = await _request_xai(
"https://cli-chat-proxy.grok.com/v1/responses",
_build_headers("secret", "grok-4.5"),
{"model": "grok-4.5", "tools": [{"type": "x_search"}]},
on_tool_call_delta=lambda event: _append(tool_events, event),
)
assert result[0] == ""
assert tool_events == [
{
"kind": "hosted_tool",
"phase": "start",
"call_id": "x-search-1",
"name": "x_search",
"arguments": {"query": "nanobot oauth"},
"result": None,
},
{
"kind": "hosted_tool",
"phase": "end",
"call_id": "x-search-1",
"name": "x_search",
"arguments": {"query": "nanobot oauth"},
"result": {"name": "x_semantic_search"},
},
]
assert "large hosted result" not in json.dumps(tool_events)
def test_model_capabilities_follow_upstream_aliases_and_default_to_disabled() -> None:
capabilities = _parse_xai_model_capabilities(
{
@@ -522,5 +589,5 @@ def test_large_json_error_body_redacts_camel_case_credentials_before_bounding()
assert len(detail) == 1001
async def _append(target: list[str], value: str) -> None:
async def _append(target: list[Any], value: Any) -> None:
target.append(value)
@@ -10,7 +10,7 @@ import {
export function WebSearchRun({ run, turnActive }: { run: WebSearchRunModel; turnActive: boolean }) {
const active = run.status === "running" && turnActive;
const status = run.status === "running" && !turnActive ? "done" : run.status;
const label = presentWebSearchAction(run.query, status);
const label = presentWebSearchAction(run.query, status, run.target);
return (
<>
@@ -77,6 +77,7 @@ const EXCLUDED_TOOLS = new Set([
"terminal",
"web_fetch",
"web_search",
"x_search",
"write_file",
]);
@@ -30,7 +30,7 @@ export function describeTraceLine(
const query = traceFieldFromArgs(args, ["query", "q", "text"]) || args || trimmed;
return {
kind: "search",
label: presentWebSearchAction(query, status),
label: presentWebSearchAction(query, status, name === "x_search" ? "x" : "web"),
detail: "",
};
}
@@ -5,6 +5,7 @@ import { redactActivityText, safeActivityDetail } from "./activity-text";
import { displayWebHost, formatCompactWebUrl, parseSafeActivityHttpUrl } from "./web-url";
export type WebSearchStatus = "running" | "done" | "error";
export type WebSearchTarget = "web" | "x";
export interface WebSearchSource {
title: string;
@@ -16,6 +17,7 @@ export interface WebSearchSource {
export interface WebSearchRunModel {
key: string;
query: string;
target: WebSearchTarget;
status: WebSearchStatus;
sources: WebSearchSource[];
error?: string;
@@ -49,10 +51,11 @@ export function webSearchRunsByTraceLine(
function webSearchRunFromEvent(event: ToolProgressEvent): WebSearchRunModel | null {
const name = compactToolName(toolEventName(event));
if (name !== "web_search") return null;
if (name !== "web_search" && name !== "x_search") return null;
const args = toolEventArguments(event);
const query = stringField(args, ["query", "q", "text"]);
const target: WebSearchTarget = name === "x_search" ? "x" : "web";
const status: WebSearchStatus = event.phase === "error"
? "error"
: event.phase === "end"
@@ -60,10 +63,11 @@ function webSearchRunFromEvent(event: ToolProgressEvent): WebSearchRunModel | nu
: "running";
return {
key: event.call_id ? `call:${event.call_id}` : formatToolCallTrace(event) ?? `web_search:${query}`,
key: event.call_id ? `call:${event.call_id}` : formatToolCallTrace(event) ?? `${name}:${query}`,
query,
target,
status,
sources: status === "done" ? webSearchSources(event.result) : [],
sources: status === "done" && target === "web" ? webSearchSources(event.result) : [],
error: status === "error" ? readableError(event.error) : undefined,
};
}
@@ -89,6 +93,7 @@ function presentWebSearchQuery(query: string): WebSearchQueryPresentation {
export function presentWebSearchAction(
query: string,
status: WebSearchStatus,
target: WebSearchTarget = "web",
): string {
const presentation = presentWebSearchQuery(query);
const verb = status === "error"
@@ -96,8 +101,11 @@ export function presentWebSearchAction(
: status === "running"
? "Searching"
: "Searched";
const target = [presentation.scope, presentation.query].filter(Boolean).join(" · ");
return target ? `${verb} ${target}` : verb;
const queryTarget = [presentation.scope, presentation.query].filter(Boolean).join(" · ");
if (target === "x") {
return queryTarget ? `${verb} X · ${queryTarget}` : `${verb} X`;
}
return queryTarget ? `${verb} ${queryTarget}` : verb;
}
function mergeWebSearchRun(
@@ -961,6 +961,35 @@ describe("AgentActivityCluster", () => {
expect(screen.getAllByTestId("activity-step")).toHaveLength(3);
});
it("renders hosted X search as an explicit search activity", () => {
const line = 'x_search({"query":"nanobot oauth"})';
render(
<AgentActivityCluster
messages={[{
id: "t-x-search",
role: "tool",
kind: "trace",
content: line,
traces: [line],
toolEvents: [{
phase: "end",
call_id: "x-search-1",
name: "x_search",
arguments: { query: "nanobot oauth" },
result: { name: "x_semantic_search" },
}],
createdAt: 1,
}]}
isTurnStreaming={false}
hasBodyBelow={false}
/>,
);
expect(screen.getByText("Searched X · nanobot oauth")).toBeInTheDocument();
expect(screen.queryByText(/Completed X search/i)).not.toBeInTheDocument();
expect(screen.getAllByTestId("activity-step")).toHaveLength(1);
});
it("redacts credentials from web search queries, titles, and links", () => {
const query = "release notes access_token=signed-secret";
const line = `web_search(${JSON.stringify({ query })})`;
@@ -32,6 +32,14 @@ describe("trace activity semantics", () => {
expect(describeTrace('web_search({"query":"status test"})', status).label).toBe(label);
});
it.each([
["running", "Searching X · status test"],
["done", "Searched X · status test"],
["error", "Could not search X · status test"],
] as const)("identifies hosted X search activity for %s", (status, label) => {
expect(describeTrace('x_search({"query":"status test"})', status).label).toBe(label);
});
it("never exposes URL credentials, query secrets, or private-network links", () => {
const publicResult = describeTrace(
'web_fetch({"url":"https://user:password@example.com/docs?api_key=secret#section"})',