Skip to content

Commit 3464db9

Browse files
napoleondclaude
napoleond
and
claude
committed
feat(atxp): use async polling for image, music, and video generation
Switch from synchronous blocking MCP calls to async initiate+poll pattern for slow generation tools, avoiding timeout issues and showing a spinner during generation. - image: image_create_image_async + image_get_image_async (3s poll) - music: music_create_async + music_get_async (5s poll) - video: create_video + wait_for_video (10s poll) - Extract getClient() and extractResult() from call-tool.ts for reuse Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 8d80edf commit 3464db9

6 files changed

Lines changed: 246 additions & 43 deletions

File tree

‎packages/atxp/src/call-tool.test.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,15 +21,15 @@ describe('Call Tool', () => {
2121
it('should map commands to correct tool names', () => {
2222
const toolMap: Record<string, string> = {
2323
search: 'search_search',
24-
image: 'image_create_image',
25-
music: 'music_create',
24+
image: 'image_create_image_async',
25+
music: 'music_create_async',
2626
video: 'create_video',
2727
x: 'x_live_search',
2828
};
2929

3030
expect(toolMap.search).toBe('search_search');
31-
expect(toolMap.image).toBe('image_create_image');
32-
expect(toolMap.music).toBe('music_create');
31+
expect(toolMap.image).toBe('image_create_image_async');
32+
expect(toolMap.music).toBe('music_create_async');
3333
expect(toolMap.video).toBe('create_video');
3434
expect(toolMap.x).toBe('x_live_search');
3535
});

‎packages/atxp/src/call-tool.ts‎

Lines changed: 30 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,9 @@ export interface ToolResult {
1616
content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>;
1717
}
1818

19-
export async function callTool(
20-
server: string,
21-
tool: string,
22-
args: Record<string, unknown>
23-
): Promise<string> {
19+
export type McpClient = Awaited<ReturnType<typeof atxpClient>>;
20+
21+
export async function getClient(server: string): Promise<McpClient> {
2422
const connection = getConnection();
2523

2624
if (!connection) {
@@ -29,31 +27,40 @@ export async function callTool(
2927
process.exit(1);
3028
}
3129

30+
return atxpClient({
31+
mcpServer: `https://${server}`,
32+
account: new ATXPAccount(connection),
33+
oAuthDb: getOAuthDb(),
34+
logger: getCliLogger(),
35+
});
36+
}
37+
38+
export function extractResult(result: ToolResult): string {
39+
if (result.content && result.content.length > 0) {
40+
const content = result.content[0];
41+
if (content.text) {
42+
return content.text;
43+
} else if (content.data && content.mimeType) {
44+
return `[${content.mimeType} data received - ${content.data.length} bytes base64]`;
45+
}
46+
}
47+
return JSON.stringify(result, null, 2);
48+
}
49+
50+
export async function callTool(
51+
server: string,
52+
tool: string,
53+
args: Record<string, unknown>
54+
): Promise<string> {
3255
try {
33-
const client = await atxpClient({
34-
mcpServer: `https://${server}`,
35-
account: new ATXPAccount(connection),
36-
oAuthDb: getOAuthDb(),
37-
logger: getCliLogger(),
38-
});
56+
const client = await getClient(server);
3957

4058
const result = (await client.callTool({
4159
name: tool,
4260
arguments: args,
4361
})) as ToolResult;
4462

45-
// Handle different content types
46-
if (result.content && result.content.length > 0) {
47-
const content = result.content[0];
48-
if (content.text) {
49-
return content.text;
50-
} else if (content.data && content.mimeType) {
51-
// For binary content (images, audio, video), return info about it
52-
return `[${content.mimeType} data received - ${content.data.length} bytes base64]`;
53-
}
54-
}
55-
56-
return JSON.stringify(result, null, 2);
63+
return extractResult(result);
5764
} catch (error) {
5865
const errorMessage = error instanceof Error ? error.message : String(error);
5966
console.error(chalk.red(`Error calling ${tool}: ${errorMessage}`));

‎packages/atxp/src/commands/commands.test.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -38,14 +38,14 @@ describe('Tool Commands', () => {
3838

3939
describe('image command', () => {
4040
const SERVER = 'image.mcp.atxp.ai';
41-
const TOOL = 'image_create_image';
41+
const TOOL = 'image_create_image_async';
4242

4343
it('should have correct server', () => {
4444
expect(SERVER).toBe('image.mcp.atxp.ai');
4545
});
4646

4747
it('should have correct tool name', () => {
48-
expect(TOOL).toBe('image_create_image');
48+
expect(TOOL).toBe('image_create_image_async');
4949
});
5050

5151
it('should validate prompt is required', () => {
@@ -60,14 +60,14 @@ describe('Tool Commands', () => {
6060

6161
describe('music command', () => {
6262
const SERVER = 'music.mcp.atxp.ai';
63-
const TOOL = 'music_create';
63+
const TOOL = 'music_create_async';
6464

6565
it('should have correct server', () => {
6666
expect(SERVER).toBe('music.mcp.atxp.ai');
6767
});
6868

6969
it('should have correct tool name', () => {
70-
expect(TOOL).toBe('music_create');
70+
expect(TOOL).toBe('music_create_async');
7171
});
7272

7373
it('should validate prompt is required', () => {

‎packages/atxp/src/commands/image.ts‎

Lines changed: 72 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
1-
import { callTool } from '../call-tool.js';
1+
import { getClient, extractResult, type ToolResult } from '../call-tool.js';
22
import chalk from 'chalk';
3+
import ora from 'ora';
34

45
const SERVER = 'image.mcp.atxp.ai';
5-
const TOOL = 'image_create_image';
6+
const POLL_INTERVAL_MS = 3000;
67

78
export async function imageCommand(prompt: string): Promise<void> {
89
if (!prompt || prompt.trim().length === 0) {
@@ -11,6 +12,73 @@ export async function imageCommand(prompt: string): Promise<void> {
1112
process.exit(1);
1213
}
1314

14-
const result = await callTool(SERVER, TOOL, { prompt: prompt.trim() });
15-
console.log(result);
15+
const client = await getClient(SERVER);
16+
17+
// Initiate async generation
18+
const initResult = (await client.callTool({
19+
name: 'image_create_image_async',
20+
arguments: { prompt: prompt.trim() },
21+
})) as ToolResult;
22+
23+
const initText = extractResult(initResult);
24+
let taskId: string | undefined;
25+
try {
26+
const parsed = JSON.parse(initText);
27+
taskId = parsed.taskId;
28+
} catch {
29+
// If the response isn't JSON with a taskId, it may have completed synchronously
30+
console.log(initText);
31+
return;
32+
}
33+
34+
if (!taskId) {
35+
console.log(initText);
36+
return;
37+
}
38+
39+
// Poll for completion
40+
const spinner = ora('Generating image...').start();
41+
try {
42+
while (true) {
43+
await new Promise((r) => setTimeout(r, POLL_INTERVAL_MS));
44+
45+
const pollResult = (await client.callTool({
46+
name: 'image_get_image_async',
47+
arguments: { taskId },
48+
})) as ToolResult;
49+
50+
const pollText = extractResult(pollResult);
51+
let parsed: any;
52+
try {
53+
parsed = JSON.parse(pollText);
54+
} catch {
55+
// Non-JSON response — treat as final
56+
spinner.stop();
57+
console.log(pollText);
58+
return;
59+
}
60+
61+
if (parsed.status === 'completed' || parsed.status === 'success') {
62+
spinner.succeed('Image generated');
63+
// Remove status field to show just the result
64+
delete parsed.status;
65+
console.log(JSON.stringify(parsed, null, 2));
66+
return;
67+
} else if (parsed.status === 'failed' || parsed.error) {
68+
spinner.fail('Image generation failed');
69+
console.error(chalk.red(parsed.error || JSON.stringify(parsed)));
70+
process.exit(1);
71+
}
72+
73+
// Still in progress — update spinner
74+
if (parsed.status) {
75+
spinner.text = `Generating image... (${parsed.status})`;
76+
}
77+
}
78+
} catch (error) {
79+
spinner.fail('Image generation failed');
80+
const errorMessage = error instanceof Error ? error.message : String(error);
81+
console.error(chalk.red(`Error: ${errorMessage}`));
82+
process.exit(1);
83+
}
1684
}

‎packages/atxp/src/commands/music.ts‎

Lines changed: 68 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
1-
import { callTool } from '../call-tool.js';
1+
import { getClient, extractResult, type ToolResult } from '../call-tool.js';
22
import chalk from 'chalk';
3+
import ora from 'ora';
34

45
const SERVER = 'music.mcp.atxp.ai';
5-
const TOOL = 'music_create';
6+
const POLL_INTERVAL_MS = 5000;
67

78
export async function musicCommand(prompt: string): Promise<void> {
89
if (!prompt || prompt.trim().length === 0) {
@@ -11,6 +12,69 @@ export async function musicCommand(prompt: string): Promise<void> {
1112
process.exit(1);
1213
}
1314

14-
const result = await callTool(SERVER, TOOL, { prompt: prompt.trim() });
15-
console.log(result);
15+
const client = await getClient(SERVER);
16+
17+
// Initiate async generation
18+
const initResult = (await client.callTool({
19+
name: 'music_create_async',
20+
arguments: { prompt: prompt.trim() },
21+
})) as ToolResult;
22+
23+
const initText = extractResult(initResult);
24+
let taskId: string | undefined;
25+
try {
26+
const parsed = JSON.parse(initText);
27+
taskId = parsed.taskId;
28+
} catch {
29+
console.log(initText);
30+
return;
31+
}
32+
33+
if (!taskId) {
34+
console.log(initText);
35+
return;
36+
}
37+
38+
// Poll for completion
39+
const spinner = ora('Generating music...').start();
40+
try {
41+
while (true) {
42+
await new Promise((r) => setTimeout(r, POLL_INTERVAL_MS));
43+
44+
const pollResult = (await client.callTool({
45+
name: 'music_get_async',
46+
arguments: { taskId },
47+
})) as ToolResult;
48+
49+
const pollText = extractResult(pollResult);
50+
let parsed: any;
51+
try {
52+
parsed = JSON.parse(pollText);
53+
} catch {
54+
spinner.stop();
55+
console.log(pollText);
56+
return;
57+
}
58+
59+
if (parsed.status === 'completed' || parsed.status === 'success') {
60+
spinner.succeed('Music generated');
61+
delete parsed.status;
62+
console.log(JSON.stringify(parsed, null, 2));
63+
return;
64+
} else if (parsed.status === 'failed' || parsed.error) {
65+
spinner.fail('Music generation failed');
66+
console.error(chalk.red(parsed.error || JSON.stringify(parsed)));
67+
process.exit(1);
68+
}
69+
70+
if (parsed.status) {
71+
spinner.text = `Generating music... (${parsed.status})`;
72+
}
73+
}
74+
} catch (error) {
75+
spinner.fail('Music generation failed');
76+
const errorMessage = error instanceof Error ? error.message : String(error);
77+
console.error(chalk.red(`Error: ${errorMessage}`));
78+
process.exit(1);
79+
}
1680
}

‎packages/atxp/src/commands/video.ts‎

Lines changed: 68 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
1-
import { callTool } from '../call-tool.js';
1+
import { getClient, extractResult, type ToolResult } from '../call-tool.js';
22
import chalk from 'chalk';
3+
import ora from 'ora';
34

45
const SERVER = 'video.mcp.atxp.ai';
5-
const TOOL = 'create_video';
6+
const POLL_INTERVAL_MS = 10000;
67

78
export async function videoCommand(prompt: string): Promise<void> {
89
if (!prompt || prompt.trim().length === 0) {
@@ -11,6 +12,69 @@ export async function videoCommand(prompt: string): Promise<void> {
1112
process.exit(1);
1213
}
1314

14-
const result = await callTool(SERVER, TOOL, { userPrompt: prompt.trim() });
15-
console.log(result);
15+
const client = await getClient(SERVER);
16+
17+
// Initiate generation — create_video returns a taskId
18+
const initResult = (await client.callTool({
19+
name: 'create_video',
20+
arguments: { userPrompt: prompt.trim() },
21+
})) as ToolResult;
22+
23+
const initText = extractResult(initResult);
24+
let taskId: string | undefined;
25+
try {
26+
const parsed = JSON.parse(initText);
27+
taskId = parsed.taskId;
28+
} catch {
29+
console.log(initText);
30+
return;
31+
}
32+
33+
if (!taskId) {
34+
console.log(initText);
35+
return;
36+
}
37+
38+
// Poll for completion using wait_for_video with short timeouts
39+
const spinner = ora('Generating video...').start();
40+
try {
41+
while (true) {
42+
await new Promise((r) => setTimeout(r, POLL_INTERVAL_MS));
43+
44+
const pollResult = (await client.callTool({
45+
name: 'wait_for_video',
46+
arguments: { taskId, timeoutSeconds: 30 },
47+
})) as ToolResult;
48+
49+
const pollText = extractResult(pollResult);
50+
let parsed: any;
51+
try {
52+
parsed = JSON.parse(pollText);
53+
} catch {
54+
spinner.stop();
55+
console.log(pollText);
56+
return;
57+
}
58+
59+
if (parsed.status === 'completed' || parsed.status === 'success') {
60+
spinner.succeed('Video generated');
61+
delete parsed.status;
62+
console.log(JSON.stringify(parsed, null, 2));
63+
return;
64+
} else if (parsed.status === 'failed' || parsed.error) {
65+
spinner.fail('Video generation failed');
66+
console.error(chalk.red(parsed.error || JSON.stringify(parsed)));
67+
process.exit(1);
68+
}
69+
70+
if (parsed.status) {
71+
spinner.text = `Generating video... (${parsed.status})`;
72+
}
73+
}
74+
} catch (error) {
75+
spinner.fail('Video generation failed');
76+
const errorMessage = error instanceof Error ? error.message : String(error);
77+
console.error(chalk.red(`Error: ${errorMessage}`));
78+
process.exit(1);
79+
}
1680
}

0 commit comments

Comments
 (0)