Skip to content
2 changes: 1 addition & 1 deletion src/entrypoints/cli.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ async function main(): Promise<void> {
shutdown1PEventLogging,
logForDebugging,
registerPermissionHandler(server, handler) {
server.setNotificationHandler(ChannelPermissionRequestNotificationSchema() as any, async notification =>
server.setNotificationHandler(ChannelPermissionRequestNotificationSchema(), async notification =>
handler(notification.params),
);
},
Expand Down
5 changes: 3 additions & 2 deletions src/hooks/useIdeLogging.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,10 @@ import { logEvent } from 'src/services/analytics/index.js'
import { z } from 'zod/v4'
import type { MCPServerConnection } from '../services/mcp/types.js'
import { getConnectedIdeClient } from '../utils/ide.js'
import type { AnyObjectSchema } from '@modelcontextprotocol/sdk/server/zod-compat.js'
import { lazySchema } from '../utils/lazySchema.js'

const LogEventSchema = lazySchema(() =>
const LogEventSchema: () => AnyObjectSchema = lazySchema(() =>
z.object({
method: z.literal('log_event'),
params: z.object({
Expand All @@ -27,7 +28,7 @@ export function useIdeLogging(mcpClients: MCPServerConnection[]): void {
if (ideClient) {
// Register the log event handler
ideClient.client.setNotificationHandler(
LogEventSchema() as any,
LogEventSchema(),
notification => {
const { eventName, eventData } = notification.params
logEvent(
Expand Down
5 changes: 3 additions & 2 deletions src/hooks/useIdeSelection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import type {
MCPServerConnection,
} from '../services/mcp/types.js'
import { getConnectedIdeClient } from '../utils/ide.js'
import type { AnyObjectSchema } from '@modelcontextprotocol/sdk/server/zod-compat.js'
import { lazySchema } from '../utils/lazySchema.js'
export type SelectionPoint = {
line: number
Expand All @@ -29,7 +30,7 @@ export type IDESelection = {
}

// Define the selection changed notification schema
const SelectionChangedSchema = lazySchema(() =>
const SelectionChangedSchema: () => AnyObjectSchema = lazySchema(() =>
z.object({
method: z.literal('selection_changed'),
params: z.object({
Expand Down Expand Up @@ -110,7 +111,7 @@ export function useIdeSelection(

// Register notification handler for selection_changed events
ideClient.client.setNotificationHandler(
SelectionChangedSchema() as any,
SelectionChangedSchema(),
notification => {
if (currentIDERef.current !== ideClient) {
return
Expand Down
5 changes: 3 additions & 2 deletions src/hooks/usePromptsFromClaudeInChrome.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,12 @@ import { callIdeRpc } from '../services/mcp/client.js';
import type { ConnectedMCPServer, MCPServerConnection } from '../services/mcp/types.js';
import type { PermissionMode } from '../types/permissions.js';
import { CLAUDE_IN_CHROME_MCP_SERVER_NAME, isTrackedClaudeInChromeTabId } from '../utils/claudeInChrome/common.js';
import type { AnyObjectSchema } from '@modelcontextprotocol/sdk/server/zod-compat.js';
import { lazySchema } from '../utils/lazySchema.js';
import { enqueuePendingNotification } from '../utils/messageQueueManager.js';

// Schema for the prompt notification from Chrome extension (JSON-RPC 2.0 format)
const ClaudeInChromePromptNotificationSchema = lazySchema(() =>
const ClaudeInChromePromptNotificationSchema: () => AnyObjectSchema = lazySchema(() =>
z.object({
method: z.literal('notifications/message'),
params: z.object({
Expand Down Expand Up @@ -48,7 +49,7 @@ export function usePromptsFromClaudeInChrome(
}

if (mcpClient) {
mcpClient.client.setNotificationHandler(ClaudeInChromePromptNotificationSchema() as any, notification => {
mcpClient.client.setNotificationHandler(ClaudeInChromePromptNotificationSchema(), notification => {
if (mcpClientRef.current !== mcpClient) {
return;
}
Expand Down
34 changes: 18 additions & 16 deletions src/services/mcp/channelNotification.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
*/

import type { ServerCapabilities } from '@modelcontextprotocol/sdk/types.js'
import type { AnyObjectSchema } from '@modelcontextprotocol/sdk/server/zod-compat.js'
import { z } from 'zod/v4'
import { type ChannelEntry, getAllowedChannels } from '../../bootstrap/state.js'
import { CHANNEL_TAG } from '../../constants/xml.js'
Expand Down Expand Up @@ -96,23 +97,24 @@ export type ChannelPermissionRequestParams = {
}
}

export const ChannelPermissionRequestNotificationSchema = lazySchema(() =>
z.object({
method: z.literal(CHANNEL_PERMISSION_REQUEST_METHOD),
params: z.object({
request_id: z.string(),
tool_name: z.string(),
description: z.string(),
input_preview: z.string(),
channel_context: z
.object({
source_server: z.string().optional(),
chat_id: z.string().optional(),
})
.optional(),
export const ChannelPermissionRequestNotificationSchema: () => AnyObjectSchema =
lazySchema(() =>
z.object({
method: z.literal(CHANNEL_PERMISSION_REQUEST_METHOD),
params: z.object({
request_id: z.string(),
tool_name: z.string(),
description: z.string(),
input_preview: z.string(),
channel_context: z
.object({
source_server: z.string().optional(),
chat_id: z.string().optional(),
})
.optional(),
}),
}),
}),
)
)

/**
* Meta keys become XML attribute NAMES — a crafted key like
Expand Down
35 changes: 25 additions & 10 deletions src/utils/forkedAgent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,14 @@ import {
} from '../services/analytics/index.js'
import { accumulateUsage, updateUsage } from '../services/api/claude.js'
import { EMPTY_USAGE, type NonNullableUsage } from '@ant/model-provider'
import type {
BetaRawMessageDeltaEvent,
BetaRawMessageStreamEvent,
} from '@anthropic-ai/sdk/resources/beta/messages/messages.js'
import type { ToolUseContext } from '../Tool.js'
import type { AgentDefinition } from '@claude-code-best/builtin-tools/tools/AgentTool/loadAgentsDir.js'
import type { AgentId } from '../types/ids.js'
import type { Message } from '../types/message.js'
import type { Message, StreamEvent } from '../types/message.js'
import { createChildAbortController } from './abortController.js'
import { logForDebugging } from './debug.js'
import { cloneFileStateCache } from './fileStateCache.js'
Expand Down Expand Up @@ -492,6 +496,24 @@ export function createSubagentContext(
* })
* ```
*/

type StreamEventMessage = StreamEvent & {
type: 'stream_event'
event: BetaRawMessageStreamEvent
}

function isMessageDeltaStreamEvent(
message: Message | StreamEvent,
): message is StreamEventMessage & { event: BetaRawMessageDeltaEvent } {
return (
message.type === 'stream_event' &&
typeof (message as StreamEventMessage).event === 'object' &&
(message as StreamEventMessage).event !== null &&
'type' in (message as StreamEventMessage).event &&
(message as StreamEventMessage).event.type === 'message_delta'
)
}

export async function runForkedAgent({
promptMessages,
cacheSafeParams,
Expand Down Expand Up @@ -562,15 +584,8 @@ export async function runForkedAgent({
})) {
// Extract real usage from message_delta stream events (final usage per API call)
if (message.type === 'stream_event') {
if (
'event' in message &&
(message as any).event?.type === 'message_delta' &&
(message as any).event.usage
) {
const turnUsage = updateUsage(
{ ...EMPTY_USAGE },
(message as any).event.usage,
)
if (isMessageDeltaStreamEvent(message)) {
const turnUsage = updateUsage({ ...EMPTY_USAGE }, message.event.usage)
totalUsage = accumulateUsage(totalUsage, turnUsage)
}
continue
Expand Down
34 changes: 26 additions & 8 deletions src/utils/hooks/execAgentHook.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,12 @@ import { type Tool, toolMatchesName } from '../../Tool.js'
import { SYNTHETIC_OUTPUT_TOOL_NAME } from '@claude-code-best/builtin-tools/tools/SyntheticOutputTool/SyntheticOutputTool.js'
import { ALL_AGENT_DISALLOWED_TOOLS } from '../../tools.js'
import { asAgentId } from '../../types/ids.js'
import type { Message } from '../../types/message.js'
import type {
AttachmentMessage,
Message,
RequestStartEvent,
StreamEvent,
} from '../../types/message.js'
import { createAbortController } from '../abortController.js'
import { createAttachmentMessage } from '../attachments.js'
import { createCombinedAbortSignal } from '../combinedAbortSignal.js'
Expand All @@ -30,6 +35,24 @@ import {
} from './hookHelpers.js'
import { clearSessionHooks } from './sessionHooks.js'

type QueryMessage = Message | StreamEvent | RequestStartEvent

type StructuredOutputAttachment = {
type: 'structured_output'
data: unknown
[key: string]: unknown
}

type StructuredOutputAttachmentMessage =
AttachmentMessage<StructuredOutputAttachment>

function isStructuredOutputAttachmentMessage(
message: QueryMessage,
): message is StructuredOutputAttachmentMessage {
if (message.type !== 'attachment') return false
return (message as Message).attachment?.type === 'structured_output'
}

/**
* Execute an agent-based hook using a multi-turn LLM query
*/
Expand Down Expand Up @@ -209,13 +232,8 @@ When done, return your result using the ${SYNTHETIC_OUTPUT_TOOL_NAME} tool with:
}

// Check for structured output in attachments
if (
message.type === 'attachment' &&
(message as any).attachment.type === 'structured_output'
) {
const parsed = hookResponseSchema().safeParse(
(message as any).attachment.data,
)
if (isStructuredOutputAttachmentMessage(message)) {
const parsed = hookResponseSchema().safeParse(message.attachment.data)
if (parsed.success) {
structuredOutputResult = parsed.data
logForDebugging(
Expand Down