feat(ai): generation middleware for media observability (#720)#760
Open
season179 wants to merge 2 commits into
Open
feat(ai): generation middleware for media observability (#720)#760season179 wants to merge 2 commits into
season179 wants to merge 2 commits into
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds a shared ChangesActivity-Agnostic GenerationMiddleware for Media Activities
Sequence Diagram(s)sequenceDiagram
participant Caller
participant runGenerateImage
participant GenerationMiddlewareRunner as Generation Middleware Runner
participant otelMiddleware
participant Tracer
participant Adapter
Caller->>runGenerateImage: generateImage({ middleware: [otelMiddleware] })
runGenerateImage->>GenerationMiddlewareRunner: runGenerationStart(middleware, mwCtx)
GenerationMiddlewareRunner->>otelMiddleware: onStart(mwCtx)
otelMiddleware->>Tracer: startSpan("image_generation", CLIENT)
otelMiddleware->>otelMiddleware: mediaSpans.set(mwCtx, span)
runGenerateImage->>Adapter: invoke image adapter
Adapter-->>runGenerateImage: result + usage
runGenerateImage->>GenerationMiddlewareRunner: runGenerationUsage(middleware, mwCtx, usage)
GenerationMiddlewareRunner->>otelMiddleware: onUsage(mwCtx, usage)
otelMiddleware->>otelMiddleware: span.setAttributes(usageAttributes(usage))
runGenerateImage->>GenerationMiddlewareRunner: runGenerationFinish(middleware, mwCtx, {durationMs, usage})
GenerationMiddlewareRunner->>otelMiddleware: onFinish(mwCtx, info)
otelMiddleware->>Tracer: span.end()
runGenerateImage-->>Caller: return result
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/advanced/otel.md`:
- Line 194: The openaiImage adapter example in the documentation is using an
outdated model identifier. In the openaiImage function call, replace the model
parameter from 'gpt-image-1' to 'gpt-image-2' to reflect the latest OpenAI image
generation model as defined in the adapter's model-meta.ts file.
In `@packages/ai/src/activities/generateVideo/index.ts`:
- Around line 375-385: The observer lifecycle is mismanaged: notifyObserverStart
at lines 375 and around 395-397 occurs after yields, allowing early breaks to
emit onError without onStart, and notifyObserverFinish at lines 459-472 occurs
after yielding generation:result, causing the finally block (lines 518-539) to
see settled === false and emit a synthetic cancellation error. Move both
notifyObserverStart calls to execute before any yields in the function, and move
notifyObserverFinish to execute before the generation:result yield (or set
settled === true before yielding), so that the finally block's settled flag
accurately reflects whether the observer lifecycle was properly started and
completed.
In `@testing/e2e/src/routes/api.otel-media.ts`:
- Around line 117-127: The payload validation is missing before destructuring
properties from the data object. Currently, if body.forwardedProps, body.data,
or body is null or a primitive value, or if required fields (prompt, provider)
are missing, the destructuring assignment will throw an unhandled error
resulting in a 500 response instead of the proper error response. Add validation
logic before the destructuring of prompt, provider, testId, and aimockPort from
the data object to check that data is a valid object with the required
properties, and return { ok: false, error } with an appropriate error message if
validation fails. Only proceed to createImageAdapter and
createLocalCaptureTracer after successful validation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 3e1da7aa-1e1b-402b-98f8-1b412af95984
📒 Files selected for processing (27)
.changeset/activity-observers.md.changeset/otel-full-usage-emission.mddocs/advanced/otel.mddocs/config.jsonpackages/ai/package.jsonpackages/ai/src/activities/generateAudio/index.tspackages/ai/src/activities/generateImage/index.tspackages/ai/src/activities/generateSpeech/index.tspackages/ai/src/activities/generateTranscription/index.tspackages/ai/src/activities/generateVideo/index.tspackages/ai/src/index.tspackages/ai/src/middlewares/otel.tspackages/ai/src/observability/index.tspackages/ai/src/observability/notify.tspackages/ai/src/observability/otel.tspackages/ai/src/observability/types.tspackages/ai/src/observability/usage-attributes.tspackages/ai/src/utilities/errors.tspackages/ai/src/utilities/numbers.tspackages/ai/tests/middlewares/otel.test.tspackages/ai/tests/observability/activity-observers.test.tspackages/ai/tests/observability/otel-observer.test.tspackages/ai/vite.config.tstesting/e2e/src/routeTree.gen.tstesting/e2e/src/routes/api.otel-media.tstesting/e2e/src/routes/api.otel-usage.tstesting/e2e/tests/middleware.spec.ts
Contributor
|
I like the idea of adding this but not like an observer Prop, chat has ChatMiddleware, this should just be GenerationMiddleware passed in via middleware prop same as chat, the chat one would be a superset that just observes more hooks than the generation one. So that way the API is uniform across both |
Contributor
|
Thanks heaps for this @season179. Would you be able to resolve conflicts first before I review? |
|
View your CI Pipeline Execution ↗ for commit a1f5bf9 ☁️ Nx Cloud last updated this comment at |
Contributor
Author
|
Thank you, Alem for the feedback. I'll make the change accordingly. @tombeckenham, yes, I'll make it conflict-free after addressing Alem's comment. |
season179
changed the title
season179
force-pushed
the
feat/activity-observers
branch
from
e835e46 to
fefbfc7
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
packages/ai/tests/middlewares/fake-otel.ts (1)
242-275:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winAvoid casting an incomplete fake context to
ChatMiddlewareContext.
makeCtxcurrently relies onas ChatMiddlewareContextwhile omitting required context fields (e.g.,runId,threadId, capability accessors). This can mask contract drift and make tests fail later for the wrong reason.Proposed fix (build a complete typed context)
import { SpanStatusCode, trace as otelTrace } from '`@opentelemetry/api`' import type { ChatMiddlewareContext } from '../../src/activities/chat/middleware/types' +import { CapabilityRegistry } from '../../src/activities/chat/middleware/capabilities' import type { ToolCall } from '../../src/types' @@ export function makeCtx( overrides: Partial<ChatMiddlewareContext> = {}, ): ChatMiddlewareContext { - const base = { + const capabilities = new CapabilityRegistry() + const base: ChatMiddlewareContext = { requestId: 'req-1', streamId: 'stream-1', + runId: 'run-1', + threadId: 'thread-1', // Chat contexts always report `activity: 'chat'`; otelMiddleware branches // on it to choose the chat span tree over the single media span. activity: 'chat' as const, @@ messages: [], createId: (prefix: string) => `${prefix}-1`, + capabilities, + get: (capability) => capability[0](base), + getOptional: (capability) => capability[0](base, { optional: true }), + provide: (capability, value) => capability[1](base, value), } - return { - ...base, - ...overrides, - } as ChatMiddlewareContext + return { ...base, ...overrides } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/ai/tests/middlewares/fake-otel.ts` around lines 242 - 275, The makeCtx function casts an incomplete object to ChatMiddlewareContext using "as ChatMiddlewareContext" to hide missing required fields, which masks contract drift and causes misleading test failures. Remove the type cast at the end of the return statement in makeCtx and instead add all required fields to the base object, including fields like runId, threadId, and any required capability accessors. This ensures the object is fully typed without relying on casting to suppress type errors.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/ai/tests/middleware-capabilities.test.ts`:
- Around line 1-17: The imports in the middleware-capabilities.test.ts file are
not ordered according to ESLint's sort-imports/import/order rules. Reorganize
the imports by grouping them in the following order: first place external
library imports (vitest) at the top, then all type imports together (type
imports from CapabilityContext, ChatMiddleware, and ChatMiddlewareContext),
followed by relative imports from ../src organized by path depth (starting with
deeper paths like ../src/activities/chat/middleware/capabilities, validate,
builder, define, compose, then ../src/logger/resolve, then ../src), and finally
place sibling imports starting with ./ (test-utils) at the end. Ensure imports
from the same source are grouped together and alphabetized within each group.
---
Outside diff comments:
In `@packages/ai/tests/middlewares/fake-otel.ts`:
- Around line 242-275: The makeCtx function casts an incomplete object to
ChatMiddlewareContext using "as ChatMiddlewareContext" to hide missing required
fields, which masks contract drift and causes misleading test failures. Remove
the type cast at the end of the return statement in makeCtx and instead add all
required fields to the base object, including fields like runId, threadId, and
any required capability accessors. This ensures the object is fully typed
without relying on casting to suppress type errors.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 7019526b-080b-4464-abed-22bf041bc3cf
📒 Files selected for processing (26)
.changeset/activity-observers.mddocs/advanced/otel.mddocs/config.jsonpackages/ai/src/activities/chat/index.tspackages/ai/src/activities/chat/middleware/types.tspackages/ai/src/activities/generateAudio/index.tspackages/ai/src/activities/generateImage/index.tspackages/ai/src/activities/generateSpeech/index.tspackages/ai/src/activities/generateTranscription/index.tspackages/ai/src/activities/generateVideo/index.tspackages/ai/src/activities/middleware/index.tspackages/ai/src/activities/middleware/run.tspackages/ai/src/activities/middleware/types.tspackages/ai/src/index.tspackages/ai/src/middlewares/otel.tspackages/ai/src/middlewares/usage-attributes.tspackages/ai/src/utilities/errors.tspackages/ai/src/utilities/numbers.tspackages/ai/tests/middleware-capabilities.test-d.tspackages/ai/tests/middleware-capabilities.test.tspackages/ai/tests/middlewares/fake-otel.tspackages/ai/tests/middlewares/generation-middleware.test.tspackages/ai/tests/middlewares/otel.test.tstesting/e2e/src/routeTree.gen.tstesting/e2e/src/routes/api.otel-media.tstesting/e2e/tests/middleware.spec.ts
✅ Files skipped from review due to trivial changes (4)
- packages/ai/src/utilities/numbers.ts
- .changeset/activity-observers.md
- docs/config.json
- docs/advanced/otel.md
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/ai/src/utilities/errors.ts
Contributor
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
packages/ai/tests/middlewares/fake-otel.ts (1)
242-275:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winAvoid casting an incomplete fake context to
ChatMiddlewareContext.
makeCtxcurrently relies onas ChatMiddlewareContextwhile omitting required context fields (e.g.,runId,threadId, capability accessors). This can mask contract drift and make tests fail later for the wrong reason.Proposed fix (build a complete typed context)
import { SpanStatusCode, trace as otelTrace } from '`@opentelemetry/api`' import type { ChatMiddlewareContext } from '../../src/activities/chat/middleware/types' +import { CapabilityRegistry } from '../../src/activities/chat/middleware/capabilities' import type { ToolCall } from '../../src/types' @@ export function makeCtx( overrides: Partial<ChatMiddlewareContext> = {}, ): ChatMiddlewareContext { - const base = { + const capabilities = new CapabilityRegistry() + const base: ChatMiddlewareContext = { requestId: 'req-1', streamId: 'stream-1', + runId: 'run-1', + threadId: 'thread-1', // Chat contexts always report `activity: 'chat'`; otelMiddleware branches // on it to choose the chat span tree over the single media span. activity: 'chat' as const, @@ messages: [], createId: (prefix: string) => `${prefix}-1`, + capabilities, + get: (capability) => capability[0](base), + getOptional: (capability) => capability[0](base, { optional: true }), + provide: (capability, value) => capability[1](base, value), } - return { - ...base, - ...overrides, - } as ChatMiddlewareContext + return { ...base, ...overrides } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/ai/tests/middlewares/fake-otel.ts` around lines 242 - 275, The makeCtx function casts an incomplete object to ChatMiddlewareContext using "as ChatMiddlewareContext" to hide missing required fields, which masks contract drift and causes misleading test failures. Remove the type cast at the end of the return statement in makeCtx and instead add all required fields to the base object, including fields like runId, threadId, and any required capability accessors. This ensures the object is fully typed without relying on casting to suppress type errors.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/ai/tests/middleware-capabilities.test.ts`:
- Around line 1-17: The imports in the middleware-capabilities.test.ts file are
not ordered according to ESLint's sort-imports/import/order rules. Reorganize
the imports by grouping them in the following order: first place external
library imports (vitest) at the top, then all type imports together (type
imports from CapabilityContext, ChatMiddleware, and ChatMiddlewareContext),
followed by relative imports from ../src organized by path depth (starting with
deeper paths like ../src/activities/chat/middleware/capabilities, validate,
builder, define, compose, then ../src/logger/resolve, then ../src), and finally
place sibling imports starting with ./ (test-utils) at the end. Ensure imports
from the same source are grouped together and alphabetized within each group.
---
Outside diff comments:
In `@packages/ai/tests/middlewares/fake-otel.ts`:
- Around line 242-275: The makeCtx function casts an incomplete object to
ChatMiddlewareContext using "as ChatMiddlewareContext" to hide missing required
fields, which masks contract drift and causes misleading test failures. Remove
the type cast at the end of the return statement in makeCtx and instead add all
required fields to the base object, including fields like runId, threadId, and
any required capability accessors. This ensures the object is fully typed
without relying on casting to suppress type errors.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 7019526b-080b-4464-abed-22bf041bc3cf
📒 Files selected for processing (26)
.changeset/activity-observers.mddocs/advanced/otel.mddocs/config.jsonpackages/ai/src/activities/chat/index.tspackages/ai/src/activities/chat/middleware/types.tspackages/ai/src/activities/generateAudio/index.tspackages/ai/src/activities/generateImage/index.tspackages/ai/src/activities/generateSpeech/index.tspackages/ai/src/activities/generateTranscription/index.tspackages/ai/src/activities/generateVideo/index.tspackages/ai/src/activities/middleware/index.tspackages/ai/src/activities/middleware/run.tspackages/ai/src/activities/middleware/types.tspackages/ai/src/index.tspackages/ai/src/middlewares/otel.tspackages/ai/src/middlewares/usage-attributes.tspackages/ai/src/utilities/errors.tspackages/ai/src/utilities/numbers.tspackages/ai/tests/middleware-capabilities.test-d.tspackages/ai/tests/middleware-capabilities.test.tspackages/ai/tests/middlewares/fake-otel.tspackages/ai/tests/middlewares/generation-middleware.test.tspackages/ai/tests/middlewares/otel.test.tstesting/e2e/src/routeTree.gen.tstesting/e2e/src/routes/api.otel-media.tstesting/e2e/tests/middleware.spec.ts
✅ Files skipped from review due to trivial changes (4)
- packages/ai/src/utilities/numbers.ts
- .changeset/activity-observers.md
- docs/config.json
- docs/advanced/otel.md
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/ai/src/utilities/errors.ts
🛑 Comments failed to post (1)
packages/ai/tests/middleware-capabilities.test.ts (1)
1-17:
⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFix import ordering to satisfy ESLint checks.
Line 1 through Line 17 currently violate
sort-imports/import/order(as flagged by static analysis), which can block CI on lint-gated pipelines.Suggested import reordering
-import { describe, it, expect, vi } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { - createCapability, CapabilityRegistry, + createCapability, } from '../src/activities/chat/middleware/capabilities' import { validateCapabilities } from '../src/activities/chat/middleware/validate' import { createChatMiddleware } from '../src/activities/chat/middleware/builder' import { defineChatMiddleware } from '../src/activities/chat/middleware/define' -import type { CapabilityContext } from '../src/activities/chat/middleware/capabilities' import { MiddlewareRunner } from '../src/activities/chat/middleware/compose' import { resolveDebugOption } from '../src/logger/resolve' import { chat } from '../src' -import { createMockAdapter, ev, collectChunks } from './test-utils' +import { collectChunks, createMockAdapter, ev } from './test-utils' +import type { CapabilityContext } from '../src/activities/chat/middleware/capabilities' import type { ChatMiddleware, ChatMiddlewareContext, } from '../src/activities/chat/middleware/types'🧰 Tools
🪛 ESLint
[error] 1-1: Member 'expect' of the import declaration should be sorted alphabetically.
(sort-imports)
[error] 4-4: Member 'CapabilityRegistry' of the import declaration should be sorted alphabetically.
(sort-imports)
[error] 9-9:
../src/activities/chat/middleware/capabilitiestype import should occur after import of./test-utils(import/order)
[error] 13-13: Member 'collectChunks' of the import declaration should be sorted alphabetically.
(sort-imports)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/ai/tests/middleware-capabilities.test.ts` around lines 1 - 17, The imports in the middleware-capabilities.test.ts file are not ordered according to ESLint's sort-imports/import/order rules. Reorganize the imports by grouping them in the following order: first place external library imports (vitest) at the top, then all type imports together (type imports from CapabilityContext, ChatMiddleware, and ChatMiddlewareContext), followed by relative imports from ../src organized by path depth (starting with deeper paths like ../src/activities/chat/middleware/capabilities, validate, builder, define, compose, then ../src/logger/resolve, then ../src), and finally place sibling imports starting with ./ (test-utils) at the end. Ensure imports from the same source are grouped together and alphabetized within each group.Source: Linters/SAST tools
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Closes #720.
Adds activity-agnostic OpenTelemetry coverage for the media generation
activities:
generateImage,generateVideo,generateAudio,generateSpeech, andgenerateTranscription.Following Alem's review feedback, this does not introduce a separate observer
prop/API. Media activities now accept the same
middlewareoption shape aschat, using a base
GenerationMiddlewarecontract.ChatMiddlewareremains thericher chat-specific superset, while observe-only middleware authored against
GenerationMiddlewarecan run across both chat and media activities.otelMiddleware()now satisfies bothGenerationMiddlewareandChatMiddleware, so the same middleware instance can be passed tochat()andto any media activity. For media calls it emits one
gen_ai.*client span peractivity, tagged with the appropriate
gen_ai.operation.name:generateImage->image_generationgenerateVideo->video_generationgenerateAudio->audio_generationgenerateSpeech->text_to_speechgenerateTranscription->transcriptionThe media spans reuse the shared usage-attribute emission from chat, including
cost, total tokens, cache/reasoning detail, duration billing, upstream cost, and
media unit counts where reported. When a
Meteris supplied, the middlewarealso records
gen_ai.client.operation.durationfor media activity calls.Streaming video is covered through the full create -> poll -> complete
lifecycle. If a consumer abandons the stream before completion, the middleware
receives
onAbortandotelMiddleware()ends the span as cancelled rather thanleaking it.
Branch status
#747 has merged into
main, and this branch has been refreshed on currentmainso the PR diff only carries the #720 media generation middleware work.Test plan
packages/ai/tests/middlewares/generation-middleware.test.tscoversmedia middleware wiring, lifecycle ordering, errors, and streaming-video abort
behavior.
packages/ai/tests/middlewares/otel.test.tscoversotelMiddlewarespan/metric behavior for chat and media generation.
packages/ai/tests/middleware-capabilities.test-d.tsverifiesotelMiddleware()satisfies bothGenerationMiddlewareandChatMiddleware,while chat-only middleware is not assignable to a media
GenerationMiddlewareslot.testing/e2eroute + spec exercisesotelMiddleware()ongenerateImagethrough the mock transport.pnpm installpnpm test:prpnpm --filter @tanstack/ai-e2e test:e2e(exit 0; 265 passed, 5 flakyafter retries)
Summary by CodeRabbit
Release Notes
CLIENTspan per media call, using activity-specific operation names, including usage/cost attributes and optional duration histograms; streaming video now spans create→poll→complete and reports abandonment/cancellation correctly.