Skip to content

Commit 576b549

Browse files
Jonesxqtombeckenham
Jonesxq
authored and
tombeckenham
committed
fix(ai-react): preserve generation devtools identity
1 parent 888e8b7 commit 576b549

10 files changed

Lines changed: 200 additions & 67 deletions

‎.changeset/react-generation-devtools-identity.md‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
---
2+
'@tanstack/ai-react': patch
3+
---
4+
5+
Prevent caller-supplied `devtools` options from overriding the framework and
6+
hook identity reported by `useGeneration`. Custom metadata continues to pass
7+
through unchanged.

‎packages/ai-react/src/use-generate-audio.ts‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { useGeneration } from './use-generation'
1+
import { useGenerationWithDevtoolsIdentity } from './use-generation'
22
import { reconstructAudioResult } from '@tanstack/ai-client'
33
import type { AudioGenerationResult, StreamChunk } from '@tanstack/ai'
44
import type {
@@ -151,21 +151,17 @@ export function useGenerateAudio<TTransformed = void>(
151151
): UseGenerateAudioReturn<
152152
InferGenerationOutputFromReturn<AudioGenerationResult, TTransformed>
153153
> {
154-
const devtools = {
155-
...options.devtools,
156-
framework: 'react',
157-
hookName: 'useGenerateAudio',
158-
outputKind: 'audio' as const,
159-
}
160-
const generation = useGeneration<
154+
const generation = useGenerationWithDevtoolsIdentity<
161155
AudioGenerateInput,
162156
AudioGenerationResult,
163157
TTransformed
164-
>({
165-
...options,
166-
devtools,
167-
reconstructResult: reconstructAudioResult,
168-
})
158+
>(
159+
{ ...options, reconstructResult: reconstructAudioResult },
160+
{
161+
hookName: 'useGenerateAudio',
162+
outputKind: 'audio',
163+
},
164+
)
169165

170166
return generation
171167
}

‎packages/ai-react/src/use-generate-image.ts‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { useGeneration } from './use-generation'
1+
import { useGenerationWithDevtoolsIdentity } from './use-generation'
22
import { reconstructImageResult } from '@tanstack/ai-client'
33
import type { ImageGenerationResult, StreamChunk } from '@tanstack/ai'
44
import type {
@@ -153,21 +153,17 @@ export function useGenerateImage<TTransformed = void>(
153153
): UseGenerateImageReturn<
154154
InferGenerationOutputFromReturn<ImageGenerationResult, TTransformed>
155155
> {
156-
const devtools = {
157-
...options.devtools,
158-
framework: 'react',
159-
hookName: 'useGenerateImage',
160-
outputKind: 'image' as const,
161-
}
162-
const generation = useGeneration<
156+
const generation = useGenerationWithDevtoolsIdentity<
163157
ImageGenerateInput,
164158
ImageGenerationResult,
165159
TTransformed
166-
>({
167-
...options,
168-
devtools,
169-
reconstructResult: reconstructImageResult,
170-
})
160+
>(
161+
{ ...options, reconstructResult: reconstructImageResult },
162+
{
163+
hookName: 'useGenerateImage',
164+
outputKind: 'image',
165+
},
166+
)
171167

172168
return generation
173169
}

‎packages/ai-react/src/use-generate-speech.ts‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { useGeneration } from './use-generation'
1+
import { useGenerationWithDevtoolsIdentity } from './use-generation'
22
import { reconstructSpeechResult } from '@tanstack/ai-client'
33
import type { StreamChunk, TTSResult } from '@tanstack/ai'
44
import type {
@@ -147,21 +147,17 @@ export function useGenerateSpeech<TTransformed = void>(
147147
): UseGenerateSpeechReturn<
148148
InferGenerationOutputFromReturn<TTSResult, TTransformed>
149149
> {
150-
const devtools = {
151-
...options.devtools,
152-
framework: 'react',
153-
hookName: 'useGenerateSpeech',
154-
outputKind: 'audio' as const,
155-
}
156-
const generation = useGeneration<
150+
const generation = useGenerationWithDevtoolsIdentity<
157151
SpeechGenerateInput,
158152
TTSResult,
159153
TTransformed
160-
>({
161-
...options,
162-
devtools,
163-
reconstructResult: reconstructSpeechResult,
164-
})
154+
>(
155+
{ ...options, reconstructResult: reconstructSpeechResult },
156+
{
157+
hookName: 'useGenerateSpeech',
158+
outputKind: 'audio',
159+
},
160+
)
165161

166162
return generation
167163
}

‎packages/ai-react/src/use-generation.ts‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { createGenerationDevtoolsBridge } from '@tanstack/ai-client/devtools'
33
import { useCallback, useEffect, useId, useMemo, useRef, useState } from 'react'
44
import type { StreamChunk } from '@tanstack/ai'
55
import type {
6+
AIDevtoolsClientMetadata,
67
AIDevtoolsDisplayOptions,
78
ConnectConnectionAdapter,
89
GenerationClientOptions,
@@ -166,6 +167,34 @@ export function useGeneration<
166167
): UseGenerationReturn<
167168
InferGenerationOutputFromReturn<TResult, TTransformed>,
168169
TInput
170+
> {
171+
return useGenerationWithDevtoolsIdentity<TInput, TResult, TTransformed>(
172+
options,
173+
{ hookName: 'useGeneration' },
174+
)
175+
}
176+
177+
interface GenerationDevtoolsIdentity {
178+
hookName: AIDevtoolsClientMetadata['hookName']
179+
outputKind?: AIDevtoolsClientMetadata['outputKind']
180+
}
181+
182+
/** @internal */
183+
export function useGenerationWithDevtoolsIdentity<
184+
TInput extends Record<string, any>,
185+
TResult,
186+
TTransformed = void,
187+
>(
188+
options: Omit<
189+
UseGenerationOptions<TInput, TResult>,
190+
'onResult' | 'persistence' | 'threadId' | 'id'
191+
> & {
192+
onResult?: (result: TResult) => TTransformed
193+
} & GenerationPersistenceOptions,
194+
devtoolsIdentity: GenerationDevtoolsIdentity,
195+
): UseGenerationReturn<
196+
InferGenerationOutputFromReturn<TResult, TTransformed>,
197+
TInput
169198
> {
170199
type TOutput = InferGenerationOutputFromReturn<TResult, TTransformed>
171200
const hookId = useId()
@@ -205,9 +234,12 @@ export function useGeneration<
205234
: {}),
206235
devtoolsBridgeFactory: createGenerationDevtoolsBridge,
207236
devtools: {
208-
hookName: 'useGeneration',
209-
framework: 'react',
210237
...opts.devtools,
238+
framework: 'react',
239+
hookName: devtoolsIdentity.hookName,
240+
...(devtoolsIdentity.outputKind !== undefined && {
241+
outputKind: devtoolsIdentity.outputKind,
242+
}),
211243
},
212244
// The transform's raw return type (`TTransformed`) and the stored output
213245
// (`TOutput`, with null/void/undefined stripped) are identical at runtime;

‎packages/ai-react/src/use-summarize.ts‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { useGeneration } from './use-generation'
1+
import { useGenerationWithDevtoolsIdentity } from './use-generation'
22
import { reconstructSummarizeResult } from '@tanstack/ai-client'
33
import type { StreamChunk, SummarizationResult } from '@tanstack/ai'
44
import type {
@@ -150,21 +150,17 @@ export function useSummarize<TTransformed = void>(
150150
): UseSummarizeReturn<
151151
InferGenerationOutputFromReturn<SummarizationResult, TTransformed>
152152
> {
153-
const devtools = {
154-
...options.devtools,
155-
framework: 'react',
156-
hookName: 'useSummarize',
157-
outputKind: 'text' as const,
158-
}
159-
const generation = useGeneration<
153+
const generation = useGenerationWithDevtoolsIdentity<
160154
SummarizeGenerateInput,
161155
SummarizationResult,
162156
TTransformed
163-
>({
164-
...options,
165-
devtools,
166-
reconstructResult: reconstructSummarizeResult,
167-
})
157+
>(
158+
{ ...options, reconstructResult: reconstructSummarizeResult },
159+
{
160+
hookName: 'useSummarize',
161+
outputKind: 'text',
162+
},
163+
)
168164

169165
return generation
170166
}

‎packages/ai-react/src/use-transcription.ts‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { useGeneration } from './use-generation'
1+
import { useGenerationWithDevtoolsIdentity } from './use-generation'
22
import { reconstructTranscriptionResult } from '@tanstack/ai-client'
33
import type { StreamChunk, TranscriptionResult } from '@tanstack/ai'
44
import type {
@@ -155,17 +155,17 @@ export function useTranscription<TTransformed = void>(
155155
): UseTranscriptionReturn<
156156
InferGenerationOutputFromReturn<TranscriptionResult, TTransformed>
157157
> {
158-
const devtools = {
159-
...options.devtools,
160-
framework: 'react',
161-
hookName: 'useTranscription',
162-
outputKind: 'text' as const,
163-
}
164-
const generation = useGeneration<
158+
const generation = useGenerationWithDevtoolsIdentity<
165159
TranscriptionGenerateInput,
166160
TranscriptionResult,
167161
TTransformed
168-
>({ ...options, devtools, reconstructResult: reconstructTranscriptionResult })
162+
>(
163+
{ ...options, reconstructResult: reconstructTranscriptionResult },
164+
{
165+
hookName: 'useTranscription',
166+
outputKind: 'text',
167+
},
168+
)
169169

170170
return generation
171171
}

‎packages/ai-react/tests/devtools-identification.test.ts‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
import { renderHook } from '@testing-library/react'
2+
import { beforeEach, describe, expect, it, vi } from 'vitest'
3+
import { useGenerateImage } from '../src/use-generate-image'
4+
import { useGeneration } from '../src/use-generation'
5+
6+
const { captureMetadata } = vi.hoisted(() => ({
7+
captureMetadata: vi.fn(),
8+
}))
9+
10+
vi.mock('@tanstack/ai-client/devtools', async (importOriginal) => {
11+
const actual =
12+
await importOriginal<typeof import('@tanstack/ai-client/devtools')>()
13+
const createGenerationDevtoolsBridge: typeof actual.createGenerationDevtoolsBridge =
14+
(options) => {
15+
captureMetadata(options.metadata)
16+
return actual.createGenerationDevtoolsBridge(options)
17+
}
18+
19+
return {
20+
...actual,
21+
createGenerationDevtoolsBridge,
22+
}
23+
})
24+
25+
describe('React generation devtools identification', () => {
26+
beforeEach(() => {
27+
captureMetadata.mockClear()
28+
})
29+
30+
it('keeps useGeneration identity while preserving caller metadata', () => {
31+
const devtools = {
32+
name: 'Custom generation',
33+
framework: 'vue',
34+
hookName: 'somethingElse',
35+
outputKind: 'image' as const,
36+
}
37+
38+
renderHook(() =>
39+
useGeneration({
40+
fetcher: async () => ({ id: 'result-1' }),
41+
devtools,
42+
}),
43+
)
44+
45+
expect(captureMetadata).toHaveBeenLastCalledWith({
46+
name: 'Custom generation',
47+
outputKind: 'image',
48+
framework: 'react',
49+
hookName: 'useGeneration',
50+
})
51+
})
52+
53+
it('keeps specialized hook identity', () => {
54+
renderHook(() =>
55+
useGenerateImage({
56+
fetcher: async () => ({
57+
id: 'image-1',
58+
images: [],
59+
model: 'test-model',
60+
}),
61+
devtools: { name: 'Image Studio' },
62+
}),
63+
)
64+
65+
expect(captureMetadata).toHaveBeenLastCalledWith({
66+
name: 'Image Studio',
67+
framework: 'react',
68+
hookName: 'useGenerateImage',
69+
outputKind: 'image',
70+
})
71+
})
72+
})

‎testing/e2e/src/routes/devtools-generation-hooks.tsx‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
import { createFileRoute } from '@tanstack/react-router'
2+
import { useState } from 'react'
23
import {
34
fetchServerSentEvents,
45
useGenerateAudio,
56
useGenerateImage,
67
useGenerateSpeech,
78
useGenerateVideo,
9+
useGeneration,
810
useSummarize,
911
useTranscription,
1012
} from '@tanstack/ai-react'
@@ -28,10 +30,17 @@ const SPEECH_TEXT = 'welcome to the guitar store'
2830
const SUMMARY_TEXT =
2931
'[summarize] The Fender Stratocaster is a versatile electric guitar'
3032
const VIDEO_PROMPT = 'a guitar being played in a store'
33+
const CUSTOM_GENERATION_DEVTOOLS = {
34+
name: 'Custom Generation',
35+
framework: 'vue',
36+
hookName: 'somethingElse',
37+
outputKind: 'text' as const,
38+
}
3139

3240
function DevtoolsGenerationHooksRoute() {
3341
const { testId, aimockPort } = Route.useSearch()
3442
const sharedBody = { testId, aimockPort }
43+
const [customGenerationMounted, setCustomGenerationMounted] = useState(false)
3544

3645
const image = useGenerateImage({
3746
id: 'generation-hooks:useGenerateImage',
@@ -226,7 +235,16 @@ function DevtoolsGenerationHooksRoute() {
226235
>
227236
Run All
228237
</button>
238+
<button
239+
type="button"
240+
data-testid="mount-custom-generation"
241+
className="rounded border border-gray-700 px-3 py-2 text-sm text-gray-200"
242+
onClick={() => setCustomGenerationMounted(true)}
243+
>
244+
Mount Custom Generation
245+
</button>
229246
</div>
247+
{customGenerationMounted ? <CustomGenerationIdentityProbe /> : null}
230248
<div className="grid gap-3 md:grid-cols-2">
231249
{hooks.map((hook) => (
232250
<section
@@ -301,6 +319,15 @@ function DevtoolsGenerationHooksRoute() {
301319
)
302320
}
303321

322+
function CustomGenerationIdentityProbe() {
323+
useGeneration({
324+
id: 'generation-hooks:useGeneration',
325+
fetcher: async () => ({ text: 'custom result' }),
326+
devtools: CUSTOM_GENERATION_DEVTOOLS,
327+
})
328+
return null
329+
}
330+
304331
function duplicateSingleImageResult(
305332
result: ImageGenerationResult,
306333
): ImageGenerationResult {

0 commit comments

Comments
 (0)