This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new e9a89e25eb0 Fix Common AI's AI Model tab for tasks inside a task loop
(#75518)
e9a89e25eb0 is described below
commit e9a89e25eb0c04d32778d60e83377705d81b979f
Author: Kaxil Naik <[email protected]>
AuthorDate: Fri Oct 23 20:58:37 2026 +0100
Fix Common AI's AI Model tab for tasks inside a task loop (#75518)
The AI Model tab read the task's model name and usage XComs by task id
and map index only. For a task inside a task loop the XCom route needs
the loop pass too, so it answered 400 "Select a region and index for
this task instance" and the tab showed no model info.
The tab now forwards the selected task instance's region_id and
region_index from the host, the same way the HITL review plugin does.
Hosts that predate the taskInstance prop keep the map-only reads. The
region helpers move to a side-effect-free module both plugins share.
---
.../providers/common/ai/plugins/www/package.json | 1 +
.../providers/common/ai/plugins/www/pnpm-lock.yaml | 131 +++++++++++++++++----
.../providers/common/ai/plugins/www/src/api.ts | 7 +-
.../ai/plugins/www/src/components/ChatPage.tsx | 2 +-
.../ai/plugins/www/src/components/ModelPanel.tsx | 24 ++--
.../ai/plugins/www/src/hooks/useModelInfo.ts | 44 ++++---
.../common/ai/plugins/www/src/hooks/useSession.ts | 3 +-
.../providers/common/ai/plugins/www/src/main.tsx | 9 +-
.../common/ai/plugins/www/src/model-api.ts | 17 ++-
.../providers/common/ai/plugins/www/src/model.tsx | 38 ++++--
.../providers/common/ai/plugins/www/src/region.ts | 49 ++++++++
.../common/ai/plugins/www/tests/main.test.mjs | 2 +-
.../common/ai/plugins/www/tests/model-api.test.ts | 68 +++++++++++
.../common/ai/plugins/www/tests/model.test.mjs | 121 +++++++++++++++++++
.../common/ai/plugins/www/tests/region.test.ts | 45 +++++++
15 files changed, 482 insertions(+), 79 deletions(-)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/package.json
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/package.json
index 495366ed9c2..c55dd72f992 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/package.json
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/package.json
@@ -41,6 +41,7 @@
"@types/react": "^19.2.14",
"@types/react-dom": "^19.2.3",
"@vitejs/plugin-react-swc": "^4.2.3",
+ "happy-dom": "^21.0.1",
"typescript": "~5.9.3",
"vite": "^8.1.6",
"vite-plugin-css-injected-by-js": "^3.5.2",
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/pnpm-lock.yaml
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/pnpm-lock.yaml
index 99020d1aba8..f5d4ef00df0 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/pnpm-lock.yaml
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/pnpm-lock.yaml
@@ -52,19 +52,22 @@ importers:
version: 19.2.3(@types/[email protected])
'@vitejs/plugin-react-swc':
specifier: ^4.2.3
- version:
4.2.3(@swc/[email protected])([email protected]([email protected])([email protected]))
+ version:
4.2.3(@swc/[email protected])([email protected](@types/[email protected])([email protected])([email protected]))
+ happy-dom:
+ specifier: ^21.0.1
+ version: 21.0.4
typescript:
specifier: ~5.9.3
version: 5.9.3
vite:
specifier: ^8.1.6
- version: 8.1.6([email protected])([email protected])
+ version: 8.1.6(@types/[email protected])([email protected])([email protected])
vite-plugin-css-injected-by-js:
specifier: ^3.5.2
- version: 3.5.2([email protected]([email protected])([email protected]))
+ version:
3.5.2([email protected](@types/[email protected])([email protected])([email protected]))
vite-plugin-dts:
specifier: ^4.5.4
- version:
4.5.4([email protected])([email protected])([email protected]([email protected])([email protected]))
+ version:
4.5.4(@types/[email protected])([email protected])([email protected])([email protected](@types/[email protected])([email protected])([email protected]))
packages:
@@ -750,6 +753,9 @@ packages:
'@types/[email protected]':
resolution: {integrity:
sha512-GsCCIZDE/p3i96vtEqx+7dBUGXrc7zeSK3wwPHIaRThS+9OhWIXRqzs4d6k1SVU8g91DrNRWxWUGhp5KXQb2VA==}
+ '@types/[email protected]':
+ resolution: {integrity:
sha512-uMzsv9stDdofIZBx4UGh4Z9ZgXnGC9seRDscilBcTHAeOFk67MZg/k9bWVfO+bk4qZmG+l52hRnLwUiH2LnJTg==}
+
'@types/[email protected]':
resolution: {integrity:
sha512-dISoDXWWQwUquiKsyZ4Ng+HX2KsPL7LyHKHQwgGFEA3IaKac4Obd+h2a/a6waisAoepJlBcx9paWqjA8/HVjCw==}
@@ -767,6 +773,12 @@ packages:
'@types/[email protected]':
resolution: {integrity:
sha512-ko/gIFJRv177XgZsZcBwnqJN5x/Gien8qNOn0D5bQU/zAzVf9Zt3BlcUiLqhV9y4ARk0GbT3tnUiPNgnTXzc/Q==}
+ '@types/[email protected]':
+ resolution: {integrity:
sha512-c2AKvDT8ToxLIOUlN51gTiHXflsfIFisS4pO7pDPoKouJCESkhZnEy623gwP9laCy5lnLDAw1vAzu2vM2YLOrA==}
+
+ '@types/[email protected]':
+ resolution: {integrity:
sha512-67MQl+fpWKVTT1NYdnmo3U4sc/xPo/zQBncVnI74qmQa0z/b+1g6iYqNmGCPbxO+zz2aklb08a0oHfegiVd0/w==}
+
'@ungap/[email protected]':
resolution: {integrity:
sha512-WmoN8qaIAo7WTYWbAZuG8PYEhn5fkz7dZrqTBZ7dtt//lL2Gwms1IcnQ5yHqjDfX8Ft5j4YzDM23f87zBfDe9g==}
deprecated: Potential CWE-502 - Update to 1.3.1 or higher
@@ -1080,6 +1092,10 @@ packages:
resolution: {integrity:
sha512-VZznLgtwhn+Mact9tfiwx64fA9erHH/MCXEUfB/0bX/6Fz6ny5EGTXYltMocqg4xFAQZtnO3DHWWXi8RiuN7cQ==}
engines: {node: 18 || 20 || >=22}
+ [email protected]:
+ resolution: {integrity:
sha512-nEh+kZOPY1w+gcCMobZ6ETUp9WfibndnosbpwB1iJk/8Gt5ZF2bhS6+B6bPYz424KtwsR6Rflc3tCz1/ghX2dQ==}
+ engines: {node: '>=4.0'}
+
[email protected]:
resolution: {integrity:
sha512-P8BjAsXvZS+VIDUI11hHCQEv74YT67YUi5JJFNWIqL235sBmjX4+qx9Muvls5ivyNENctx46xQLQ3aTuE7ssaQ==}
engines: {node: '>=6'}
@@ -1216,6 +1232,10 @@ packages:
[email protected]:
resolution: {integrity:
sha512-RbJ5/jmFcNNCcDV5o9eTnBLJ/HszWV0P73bc+Ff4nS/rJj+YaS6IGyiOL0VoBYX+l1Wrl3k63h/KrH+nhJ0XvQ==}
+ [email protected]:
+ resolution: {integrity:
sha512-tFm+GaLBgo6o0awHnvHaNoraVo27qJ+2tZEtZ+cIa08hXbeitrxAJajLI05hsUy/2/YTKcUjQYzHe5mCcraLBw==}
+ engines: {node: '>=20.0.0'}
+
[email protected]:
resolution: {integrity:
sha512-EykJT/Q1KjTWctppgIAgfSO0tKVuZUjhgMr17kqTumMl6Afv3EISleU7qZUzoXDFTAHTDC4NOoG/ZxU3EvlMPQ==}
engines: {node: '>=8'}
@@ -1735,6 +1755,9 @@ packages:
[email protected]:
resolution: {integrity:
sha512-yDJTmhydvl5lJzBmy/hyOAA0d+aqCBuwl818haVdYCRrWV84o7YyeVm4QlVHStqNrrJSTb6jKuFAVqAFsr+K3Q==}
+ [email protected]:
+ resolution: {integrity:
sha512-KTDyRTYX8sWmKXAikPHHSyc63CRPETMctyjKFupcC6OBLXT3xsN0e9aF7m+mIXutFWpUXuedtowG7iLOzp0kQg==}
+
[email protected]:
resolution: {integrity:
sha512-xKvGhPWw3k84Qjh8bI3ZeJjqnyadK+GEFtazSfZv/rKeTkTjOJho6mFqh2SM96iIcZokxiOpg78GazTSg8+KHA==}
@@ -1826,6 +1849,22 @@ packages:
[email protected]:
resolution: {integrity:
sha512-/BpdSx+yCQGnCvecbyXdxHDkuk55/G3xwnC0GqY4gmQ3j+A+g8kzzgB4Nk/SINjqn6+waqw3EgbVF2QKExkRxQ==}
+ [email protected]:
+ resolution: {integrity:
sha512-nt+N2dzIutVRxARx1nghPKGv1xHikU7HKdfafKkLNLindmPU/ch3U31NOCGGA/dmPcmb1VlofO0vnKAcsm0o/Q==}
+ engines: {node: '>=12'}
+
+ [email protected]:
+ resolution: {integrity:
sha512-Ydggc987+RO0AnWtZ/7Wq9FtNvcrL1b/RO0ud9mWjUPgDrsAAwQSF51sm2hm1XofbU/4jkpGEsLFsZZxU+1DOg==}
+ engines: {node: '>=10.0.0'}
+ peerDependencies:
+ bufferutil: ^4.0.1
+ utf-8-validate: '>=5.0.2'
+ peerDependenciesMeta:
+ bufferutil:
+ optional: true
+ utf-8-validate:
+ optional: true
+
[email protected]:
resolution: {integrity:
sha512-3wdGidZyq5PB084XLES5TpOSRA3wjXAlIWMhum2kRcv/41Sn2emQ0dycQW4uZXLejwKvg6EsvbdlVL+FYEct7A==}
@@ -2173,23 +2212,23 @@ snapshots:
'@jridgewell/resolve-uri': 3.1.2
'@jridgewell/sourcemap-codec': 1.5.5
- '@microsoft/[email protected]':
+ '@microsoft/[email protected](@types/[email protected])':
dependencies:
'@microsoft/tsdoc': 0.16.0
'@microsoft/tsdoc-config': 0.18.1
- '@rushstack/node-core-library': 5.20.3
+ '@rushstack/node-core-library': 5.20.3(@types/[email protected])
transitivePeerDependencies:
- '@types/node'
- '@microsoft/[email protected]':
+ '@microsoft/[email protected](@types/[email protected])':
dependencies:
- '@microsoft/api-extractor-model': 7.33.4
+ '@microsoft/api-extractor-model': 7.33.4(@types/[email protected])
'@microsoft/tsdoc': 0.16.0
'@microsoft/tsdoc-config': 0.18.1
- '@rushstack/node-core-library': 5.20.3
+ '@rushstack/node-core-library': 5.20.3(@types/[email protected])
'@rushstack/rig-package': 0.7.2
- '@rushstack/terminal': 0.22.3
- '@rushstack/ts-command-line': 5.3.3
+ '@rushstack/terminal': 0.22.3(@types/[email protected])
+ '@rushstack/ts-command-line': 5.3.3(@types/[email protected])
diff: 8.0.3
lodash: 4.18.1
minimatch: 10.2.4
@@ -2356,7 +2395,7 @@ snapshots:
'@rollup/[email protected]':
optional: true
- '@rushstack/[email protected]':
+ '@rushstack/[email protected](@types/[email protected])':
dependencies:
ajv: 8.18.0
ajv-draft-04: 1.0.0([email protected])
@@ -2366,23 +2405,29 @@ snapshots:
jju: 1.4.0
resolve: 1.22.11
semver: 7.5.4
+ optionalDependencies:
+ '@types/node': 26.6.5
- '@rushstack/[email protected]': {}
+ '@rushstack/[email protected](@types/[email protected])':
+ optionalDependencies:
+ '@types/node': 26.6.5
'@rushstack/[email protected]':
dependencies:
resolve: 1.22.11
strip-json-comments: 3.1.1
- '@rushstack/[email protected]':
+ '@rushstack/[email protected](@types/[email protected])':
dependencies:
- '@rushstack/node-core-library': 5.20.3
- '@rushstack/problem-matcher': 0.2.1
+ '@rushstack/node-core-library': 5.20.3(@types/[email protected])
+ '@rushstack/problem-matcher': 0.2.1(@types/[email protected])
supports-color: 8.1.1
+ optionalDependencies:
+ '@types/node': 26.6.5
- '@rushstack/[email protected]':
+ '@rushstack/[email protected](@types/[email protected])':
dependencies:
- '@rushstack/terminal': 0.22.3
+ '@rushstack/terminal': 0.22.3(@types/[email protected])
'@types/argparse': 1.0.38
argparse: 1.0.10
string-argv: 0.3.2
@@ -2473,6 +2518,10 @@ snapshots:
'@types/[email protected]': {}
+ '@types/[email protected]':
+ dependencies:
+ undici-types: 8.9.0
+
'@types/[email protected]': {}
'@types/[email protected](@types/[email protected])':
@@ -2487,13 +2536,19 @@ snapshots:
'@types/[email protected]': {}
+ '@types/[email protected]': {}
+
+ '@types/[email protected]':
+ dependencies:
+ '@types/node': 26.6.5
+
'@ungap/[email protected]': {}
-
'@vitejs/[email protected](@swc/[email protected])([email protected]([email protected])([email protected]))':
+
'@vitejs/[email protected](@swc/[email protected])([email protected](@types/[email protected])([email protected])([email protected]))':
dependencies:
'@rolldown/pluginutils': 1.0.0-rc.2
'@swc/core': 1.15.18(@swc/[email protected])
- vite: 8.1.6([email protected])([email protected])
+ vite: 8.1.6(@types/[email protected])([email protected])([email protected])
transitivePeerDependencies:
- '@swc/helpers'
@@ -3135,6 +3190,10 @@ snapshots:
dependencies:
balanced-match: 4.0.4
+ [email protected]:
+ dependencies:
+ '@types/node': 26.6.5
+
[email protected]: {}
[email protected]: {}
@@ -3258,6 +3317,19 @@ snapshots:
[email protected]: {}
+ [email protected]:
+ dependencies:
+ '@types/node': 26.6.5
+ '@types/whatwg-mimetype': 3.0.2
+ '@types/ws': 8.18.2
+ buffer-image-size: 0.6.4
+ entities: 7.0.1
+ whatwg-mimetype: 3.0.0
+ ws: 8.22.0
+ transitivePeerDependencies:
+ - bufferutil
+ - utf-8-validate
+
[email protected]: {}
[email protected]:
@@ -4027,6 +4099,8 @@ snapshots:
[email protected]: {}
+ [email protected]: {}
+
[email protected]:
dependencies:
'@types/unist': 3.0.3
@@ -4074,13 +4148,13 @@ snapshots:
'@types/unist': 3.0.3
vfile-message: 4.0.3
- [email protected]([email protected]([email protected])([email protected])):
+
[email protected]([email protected](@types/[email protected])([email protected])([email protected])):
dependencies:
- vite: 8.1.6([email protected])([email protected])
+ vite: 8.1.6(@types/[email protected])([email protected])([email protected])
-
[email protected]([email protected])([email protected])([email protected]([email protected])([email protected])):
+
[email protected](@types/[email protected])([email protected])([email protected])([email protected](@types/[email protected])([email protected])([email protected])):
dependencies:
- '@microsoft/api-extractor': 7.57.6
+ '@microsoft/api-extractor': 7.57.6(@types/[email protected])
'@rollup/pluginutils': 5.3.0([email protected])
'@volar/typescript': 2.4.28
'@vue/language-core': 2.2.0([email protected])
@@ -4091,13 +4165,13 @@ snapshots:
magic-string: 0.30.21
typescript: 5.9.3
optionalDependencies:
- vite: 8.1.6([email protected])([email protected])
+ vite: 8.1.6(@types/[email protected])([email protected])([email protected])
transitivePeerDependencies:
- '@types/node'
- rollup
- supports-color
- [email protected]([email protected])([email protected]):
+ [email protected](@types/[email protected])([email protected])([email protected]):
dependencies:
lightningcss: 1.33.0
picomatch: 4.0.7
@@ -4105,12 +4179,17 @@ snapshots:
rolldown: 1.1.5
tinyglobby: 0.2.17
optionalDependencies:
+ '@types/node': 26.6.5
esbuild: 0.27.7
fsevents: 2.3.3
yaml: 2.8.3
[email protected]: {}
+ [email protected]: {}
+
+ [email protected]: {}
+
[email protected]: {}
[email protected]: {}
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/api.ts
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/api.ts
index f04296e6a88..3200e398875 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/api.ts
+++ b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/api.ts
@@ -17,6 +17,7 @@
* under the License.
*/
+import { regionQuery, type TaskRegion } from "./region.ts";
import type { SessionResponse } from "src/types/feedback";
function getBase(): string {
@@ -29,8 +30,6 @@ function getBase(): string {
const BASE = getBase();
-export type TaskRegion = { region_id: string; region_index: number };
-
function buildQs(dagId: string, runId: string, taskId: string, mapIndex:
number, region?: TaskRegion): string {
const query = (
`dag_id=${encodeURIComponent(dagId)}` +
@@ -38,9 +37,7 @@ function buildQs(dagId: string, runId: string, taskId:
string, mapIndex: number,
`&task_id=${encodeURIComponent(taskId)}` +
`&map_index=${mapIndex}`
);
- return region === undefined
- ? query
- :
`${query}®ion_id=${encodeURIComponent(region.region_id)}®ion_index=${region.region_index}`;
+ return `${query}${regionQuery(region)}`;
}
export class ApiError extends Error {
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ChatPage.tsx
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ChatPage.tsx
index 7554548bb3b..38b57ccef4c 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ChatPage.tsx
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ChatPage.tsx
@@ -38,7 +38,7 @@ import {
useState,
} from "react";
-import type { TaskRegion } from "src/api";
+import type { TaskRegion } from "src/region";
import { MessageBubble } from "src/components/MessageBubble";
import { NoSession } from "src/components/NoSession";
import { useSession } from "src/hooks/useSession";
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ModelPanel.tsx
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ModelPanel.tsx
index 7d729679484..78941cf1b87 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ModelPanel.tsx
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ModelPanel.tsx
@@ -24,12 +24,14 @@ import { LuInfo } from "react-icons/lu";
import { NoModelInfo } from "src/components/NoModelInfo";
import { Tooltip } from "src/components/Tooltip";
import { useModelInfo } from "src/hooks/useModelInfo";
+import type { TaskRegion } from "src/region";
interface ModelPanelProps {
dagId: string;
runId: string;
taskId: string;
mapIndex: number;
+ region?: TaskRegion;
}
// Pydantic-ai prices a response via genai-prices' per-token rate lookup (see
@@ -74,18 +76,20 @@ const StatBox: FC<{ label: string; tooltip?: ReactNode;
value: number | string }
</Box>
);
-export const ModelPanel: FC<ModelPanelProps> = ({ dagId, runId, taskId,
mapIndex }) => {
- const { modelName, usage, loading, error } = useModelInfo(dagId, runId,
taskId, mapIndex);
+export const ModelLoading: FC = () => (
+ <Flex align="center" gap={2} p={2}>
+ <Spinner colorPalette="brand" size="sm" />
+ <Text color="fg.muted" fontSize="sm">
+ Loading model info...
+ </Text>
+ </Flex>
+);
+
+export const ModelPanel: FC<ModelPanelProps> = ({ dagId, runId, taskId,
mapIndex, region }) => {
+ const { modelName, usage, loading, error } = useModelInfo(dagId, runId,
taskId, mapIndex, region);
if (loading) {
- return (
- <Flex align="center" gap={2} p={2}>
- <Spinner colorPalette="brand" size="sm" />
- <Text color="fg.muted" fontSize="sm">
- Loading model info...
- </Text>
- </Flex>
- );
+ return <ModelLoading />;
}
if (!modelName && !usage) {
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/hooks/useModelInfo.ts
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/hooks/useModelInfo.ts
index 8b8bc24dbd8..c4b4e4df85a 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/hooks/useModelInfo.ts
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/hooks/useModelInfo.ts
@@ -17,9 +17,10 @@
* under the License.
*/
-import { useCallback, useEffect, useState } from "react";
+import { useEffect, useState } from "react";
import { createModelApi } from "src/model-api";
+import { regionOf, type TaskRegion } from "src/region";
import type { UsageInfo } from "src/types/model";
interface UseModelInfoReturn {
@@ -39,29 +40,42 @@ export function useModelInfo(
runId: string,
taskId: string,
mapIndex: number,
+ region?: TaskRegion,
): UseModelInfoReturn {
const [modelName, setModelName] = useState<string | null>(null);
const [usage, setUsage] = useState<UsageInfo | null>(null);
const [loading, setLoading] = useState(true);
const [error, setError] = useState<string | null>(null);
- // Built fresh on every call (not memoized via a ref) so a switch to another
task or map
- // index -- which changes these props without remounting the panel -- is
reflected here too.
- const fetchInfo = useCallback(async () => {
+ const regionId = region?.region_id;
+ const regionIndex = region?.region_index;
+
+ // Keyed on the region's fields, not the object, which the caller rebuilds
on every render.
+ // The host keeps the panel mounted when only the loop pass (a search param)
changes, so a
+ // read for the previous pass can still be in flight; its result is dropped.
+ useEffect(() => {
+ let current = true;
setLoading(true);
- const api = createModelApi(dagId, runId, taskId, mapIndex);
+ const api = createModelApi(
+ dagId,
+ runId,
+ taskId,
+ mapIndex,
+ regionOf({ region_id: regionId, region_index: regionIndex }),
+ );
// allSettled, not all: one XCom missing or erroring must not hide a model
name or usage
// the other call did resolve.
- const [nameResult, usageResult] = await
Promise.allSettled([api.fetchModelName(), api.fetchUsage()]);
- setModelName(nameResult.status === "fulfilled" ? nameResult.value : null);
- setUsage(usageResult.status === "fulfilled" ? usageResult.value : null);
- setError(describeFailure(nameResult) ?? describeFailure(usageResult));
- setLoading(false);
- }, [dagId, runId, taskId, mapIndex]);
-
- useEffect(() => {
- void fetchInfo();
- }, [fetchInfo]);
+ void Promise.allSettled([api.fetchModelName(),
api.fetchUsage()]).then(([nameResult, usageResult]) => {
+ if (!current) return;
+ setModelName(nameResult.status === "fulfilled" ? nameResult.value :
null);
+ setUsage(usageResult.status === "fulfilled" ? usageResult.value : null);
+ setError(describeFailure(nameResult) ?? describeFailure(usageResult));
+ setLoading(false);
+ });
+ return () => {
+ current = false;
+ };
+ }, [dagId, runId, taskId, mapIndex, regionId, regionIndex]);
return { modelName, usage, loading, error };
}
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/hooks/useSession.ts
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/hooks/useSession.ts
index 69e4de8367e..23c993598ef 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/hooks/useSession.ts
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/hooks/useSession.ts
@@ -19,7 +19,8 @@
import { useCallback, useEffect, useRef, useState } from "react";
-import { ApiError, createApi, type TaskRegion } from "src/api";
+import { ApiError, createApi } from "src/api";
+import type { TaskRegion } from "src/region";
import { type SessionResponse, isTerminalStatus } from "src/types/feedback";
const POLL_INTERVAL_MS = 3000;
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/main.tsx
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/main.tsx
index 03193798ae0..159fd1a224c 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/main.tsx
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/main.tsx
@@ -22,6 +22,7 @@ import type { FC } from "react";
import { ChatPage } from "src/components/ChatPage";
import { NoSession } from "src/components/NoSession";
+import { type HostTaskInstance, isSelectionPending, regionOf } from
"src/region";
import { Toaster } from "src/toaster";
import { localSystem } from "./theme";
@@ -31,7 +32,7 @@ export interface PluginComponentProps {
mapIndex?: string;
runId?: string;
taskId?: string;
- taskInstance?: { id?: string; region_id?: string; region_index?: number };
+ taskInstance?: HostTaskInstance;
}
/**
@@ -51,13 +52,11 @@ const PluginComponent: FC<PluginComponentProps> = (props)
=> {
} = props;
const mapIndex = /^-?\d+$/.test(String(mapIndexProp)) ?
parseInt(String(mapIndexProp), 10) : -1;
- if (!dagId || !runId || !taskId || ("taskInstance" in props && taskInstance
== null)) {
+ if (!dagId || !runId || !taskId || isSelectionPending(props)) {
return <NoSession />;
}
- const region = taskInstance?.region_id !== undefined &&
taskInstance.region_index !== undefined
- ? { region_id: taskInstance.region_id, region_index:
taskInstance.region_index }
- : undefined;
+ const region = regionOf(taskInstance);
// Remounts ChatPage on any selection change: useSession keeps its first
createApi(...) in a ref,
// so without a new key it would keep sending the previous pass's region.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model-api.ts
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model-api.ts
index 66770627c5f..d1b1af22ba6 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model-api.ts
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model-api.ts
@@ -17,6 +17,7 @@
* under the License.
*/
+import { regionQuery, type TaskRegion } from "./region.ts";
import type { UsageInfo } from "src/types/model";
function getApiBase(): string {
@@ -40,13 +41,15 @@ async function fetchXComValue<T>(
taskId: string,
mapIndex: number,
key: string,
+ region?: TaskRegion,
): Promise<T | null> {
const path =
`${API_BASE}/dags/${encodeURIComponent(dagId)}` +
`/dagRuns/${encodeURIComponent(runId)}` +
`/taskInstances/${encodeURIComponent(taskId)}` +
`/xcomEntries/${encodeURIComponent(key)}`;
- const res = await fetch(`${path}?map_index=${mapIndex}&deserialize=true`, {
+ // A task inside a task loop needs its pass's region, or the route rejects
the read as ambiguous.
+ const res = await
fetch(`${path}?map_index=${mapIndex}&deserialize=true${regionQuery(region)}`, {
credentials: "same-origin",
});
if (res.status === 404) return null;
@@ -66,12 +69,18 @@ function isUsageInfo(value: unknown): value is UsageInfo {
return typeof value === "object" && value !== null;
}
-export function createModelApi(dagId: string, runId: string, taskId: string,
mapIndex: number) {
+export function createModelApi(
+ dagId: string,
+ runId: string,
+ taskId: string,
+ mapIndex: number,
+ region?: TaskRegion,
+) {
return {
fetchModelName: () =>
- fetchXComValue<string>(dagId, runId, taskId, mapIndex,
MODEL_NAME_XCOM_KEY),
+ fetchXComValue<string>(dagId, runId, taskId, mapIndex,
MODEL_NAME_XCOM_KEY, region),
fetchUsage: async () => {
- const value = await fetchXComValue<unknown>(dagId, runId, taskId,
mapIndex, "usage");
+ const value = await fetchXComValue<unknown>(dagId, runId, taskId,
mapIndex, "usage", region);
// A reference-storing XCom backend (e.g. common.io with a low
// xcom_objectstorage_threshold) returns the stored reference string
here instead of
// the real value, since this public endpoint never runs the backend's
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model.tsx
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model.tsx
index a5bad684625..25431e72acd 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model.tsx
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model.tsx
@@ -20,8 +20,9 @@
import { Box, ChakraProvider } from "@chakra-ui/react";
import type { FC } from "react";
-import { ModelPanel } from "src/components/ModelPanel";
+import { ModelLoading, ModelPanel } from "src/components/ModelPanel";
import { NoModelInfo } from "src/components/NoModelInfo";
+import { type HostTaskInstance, isSelectionPending, regionOf } from
"src/region";
import { localSystem } from "./theme";
@@ -30,28 +31,43 @@ export interface PluginComponentProps {
mapIndex?: string;
runId?: string;
taskId?: string;
+ taskInstance?: HostTaskInstance;
}
/**
* Main plugin component for the AI Model tab.
*
- * Receives dagId, runId, taskId, mapIndex as props from the Airflow React
plugin
- * host (from route params). Renders ModelPanel when all params are present,
- * otherwise shows NoModelInfo fallback.
+ * Receives dagId, runId, taskId, mapIndex (from route params) and the selected
+ * taskInstance as props from the Airflow React plugin host. Renders
ModelPanel when
+ * all params are present, otherwise shows NoModelInfo fallback.
*/
-const PluginComponent: FC<PluginComponentProps> = ({
- dagId = "",
- runId = "",
- taskId = "",
- mapIndex: mapIndexProp = "-1",
-}) => {
+const PluginComponent: FC<PluginComponentProps> = (props) => {
+ const {
+ dagId = "",
+ runId = "",
+ taskId = "",
+ mapIndex: mapIndexProp = "-1",
+ taskInstance,
+ } = props;
const mapIndex = /^-?\d+$/.test(String(mapIndexProp)) ?
parseInt(String(mapIndexProp), 10) : -1;
if (!dagId || !runId || !taskId) {
return <NoModelInfo />;
}
- return <ModelPanel dagId={dagId} runId={runId} taskId={taskId}
mapIndex={mapIndex} />;
+ if (isSelectionPending(props)) {
+ return <ModelLoading />;
+ }
+
+ return (
+ <ModelPanel
+ dagId={dagId}
+ runId={runId}
+ taskId={taskId}
+ mapIndex={mapIndex}
+ region={regionOf(taskInstance)}
+ />
+ );
};
/**
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/region.ts
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/region.ts
new file mode 100644
index 00000000000..326388d8b62
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/region.ts
@@ -0,0 +1,49 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+// Shared by both plugin bundles, so it must stay free of top-level side
effects (api.ts and
+// model-api.ts read `document` at import time). Those two import it as
"./region.ts" because
+// the node:test runner loads them without Vite and cannot resolve the "src/"
alias.
+
+export type TaskRegion = { region_id: string; region_index: number };
+
+/** The slice of the host's TaskInstanceResponse the plugins read. */
+export type HostTaskInstance = { id?: string; region_id?: string;
region_index?: number };
+
+/**
+ * Whether a host that passes `taskInstance` has not resolved it yet. Reading
before it
+ * resolves would address a loop task without its pass; a host that predates
`taskInstance`
+ * never passes the key and keeps the map-only requests.
+ */
+export function isSelectionPending(props: { taskInstance?: HostTaskInstance
}): boolean {
+ return "taskInstance" in props && props.taskInstance == null;
+}
+
+/** The loop pass a task instance belongs to; the host omits both fields
outside a loop. */
+export function regionOf(taskInstance?: HostTaskInstance): TaskRegion |
undefined {
+ return taskInstance?.region_id !== undefined && taskInstance.region_index
!== undefined
+ ? { region_id: taskInstance.region_id, region_index:
taskInstance.region_index }
+ : undefined;
+}
+
+export function regionQuery(region?: TaskRegion): string {
+ return region === undefined
+ ? ""
+ :
`®ion_id=${encodeURIComponent(region.region_id)}®ion_index=${region.region_index}`;
+}
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/main.test.mjs
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/main.test.mjs
index 8bb563ab251..230cc118def 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/main.test.mjs
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/main.test.mjs
@@ -24,7 +24,7 @@ import { renderToStaticMarkup } from "react-dom/server";
import { createServer } from "vite";
-const server = await createServer({ server: { middlewareMode: true } });
+const server = await createServer({ server: { middlewareMode: true, ws: false
} });
after(() => server.close());
const { default: Plugin } = await server.ssrLoadModule("src/main.tsx");
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/model-api.test.ts
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/model-api.test.ts
new file mode 100644
index 00000000000..b9602343940
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/model-api.test.ts
@@ -0,0 +1,68 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import assert from "node:assert/strict";
+import { afterEach, mock, test } from "node:test";
+
+import { createModelApi } from "../src/model-api.ts";
+
+afterEach(() => mock.restoreAll());
+
+function recordXComReads(): URL[] {
+ const calls: URL[] = [];
+ mock.method(globalThis, "fetch", async (url: string) => {
+ calls.push(new URL(url, "http://localhost";));
+ return new Response(JSON.stringify({ value: {} }), { status: 200 });
+ });
+ return calls;
+}
+
+test("both model XCom reads address the loop pass", async () => {
+ const calls = recordXComReads();
+ const region = { region_id: "00000000-0000-0000-0000-000000000123",
region_index: 2 };
+ const api = createModelApi("dag", "run", "body.agent", -1, region);
+
+ await api.fetchModelName();
+ await api.fetchUsage();
+
+ assert.deepEqual(
+ calls.map((url) => url.pathname.split("/").at(-1)),
+ ["__AIRFLOW__COMMON_AI_MODEL_NAME__", "usage"],
+ );
+ for (const url of calls) {
+ assert.equal(url.searchParams.get("region_id"), region.region_id);
+ assert.equal(url.searchParams.get("region_index"), "2");
+ assert.equal(url.searchParams.get("map_index"), "-1");
+ }
+});
+
+test("a task outside a loop keeps the map-only XCom read", async () => {
+ const calls = recordXComReads();
+
+ const api = createModelApi("dag", "run", "agent", 0);
+ await api.fetchModelName();
+ await api.fetchUsage();
+
+ assert.equal(calls.length, 2);
+ for (const url of calls) {
+ assert.equal(url.searchParams.get("map_index"), "0");
+ assert.equal(url.searchParams.has("region_id"), false);
+ assert.equal(url.searchParams.has("region_index"), false);
+ }
+});
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/model.test.mjs
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/model.test.mjs
new file mode 100644
index 00000000000..fde44e545c6
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/model.test.mjs
@@ -0,0 +1,121 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import assert from "node:assert/strict";
+import { after, afterEach, mock, test } from "node:test";
+import { act, createElement } from "react";
+import { createRoot } from "react-dom/client";
+import { renderToStaticMarkup } from "react-dom/server";
+
+import { Window } from "happy-dom";
+import { createServer } from "vite";
+
+const server = await createServer({ server: { middlewareMode: true, ws: false
} });
+after(() => server.close());
+const { default: Plugin } = await server.ssrLoadModule("src/model.tsx");
+
+const taskProps = { dagId: "dag", runId: "run", taskId: "refine.agent",
mapIndex: "-1" };
+const regionId = "00000000-0000-0000-0000-000000000123";
+
+afterEach(() => mock.restoreAll());
+
+/** Mount the plugin into a happy-dom document, as the host does, and return
its controls. */
+function mountPlugin(t) {
+ const window = new Window({ url: "http://localhost/"; });
+ globalThis.window = window;
+ globalThis.document = window.document;
+ globalThis.IS_REACT_ACT_ENVIRONMENT = true;
+ const root =
createRoot(document.body.appendChild(document.createElement("div")));
+ t.after(async () => {
+ await act(async () => root.unmount());
+ await window.happyDOM.close();
+ delete globalThis.window;
+ delete globalThis.document;
+ delete globalThis.IS_REACT_ACT_ENVIRONMENT;
+ });
+ return {
+ render: (props) => act(async () => root.render(createElement(Plugin,
props))),
+ text: () => document.body.textContent,
+ };
+}
+
+/** Answer each XCom read with the model name for its loop pass, after
`delay(url)` resolves. */
+function serveXComs(delay = () => Promise.resolve()) {
+ const calls = [];
+ mock.method(globalThis, "fetch", async (input) => {
+ const url = new URL(input, "http://localhost";);
+ calls.push(url);
+ await delay(url);
+ const key = url.pathname.split("/").at(-1);
+ const value = key === "usage" ? { requests: 1 } :
`model-for-pass-${url.searchParams.get("region_index")}`;
+ return new Response(JSON.stringify({ value }), { status: 200 });
+ });
+ return calls;
+}
+
+test("an unresolved modern host selection waits instead of reading without a
region", () => {
+ const html = renderToStaticMarkup(createElement(Plugin, { ...taskProps,
taskInstance: undefined }));
+
+ assert.match(html, /Loading model info/);
+});
+
+test("an older host without task instance context reads without a region",
async (t) => {
+ const calls = serveXComs();
+ const plugin = mountPlugin(t);
+
+ await plugin.render(taskProps);
+
+ assert.equal(calls.length, 2);
+ for (const url of calls) assert.equal(url.searchParams.has("region_id"),
false);
+});
+
+test("the selected loop pass reaches both XCom reads and a pass switch
rereads", async (t) => {
+ const calls = serveXComs();
+ const plugin = mountPlugin(t);
+
+ await plugin.render({ ...taskProps, taskInstance: { region_id: regionId,
region_index: 2 } });
+ await plugin.render({ ...taskProps, taskInstance: { region_id: regionId,
region_index: 3 } });
+
+ assert.deepEqual(
+ calls.map((url) => [url.pathname.split("/").at(-1),
url.searchParams.get("region_id"), url.searchParams.get("region_index")]),
+ [
+ ["__AIRFLOW__COMMON_AI_MODEL_NAME__", regionId, "2"],
+ ["usage", regionId, "2"],
+ ["__AIRFLOW__COMMON_AI_MODEL_NAME__", regionId, "3"],
+ ["usage", regionId, "3"],
+ ],
+ );
+ assert.match(plugin.text(), /model-for-pass-3/);
+});
+
+test("a slow read for the previous loop pass does not replace the current
pass", async (t) => {
+ let releasePass2;
+ const pass2Held = new Promise((resolve) => {
+ releasePass2 = resolve;
+ });
+ serveXComs((url) => (url.searchParams.get("region_index") === "2" ?
pass2Held : Promise.resolve()));
+ const plugin = mountPlugin(t);
+
+ await plugin.render({ ...taskProps, taskInstance: { region_id: regionId,
region_index: 2 } });
+ await plugin.render({ ...taskProps, taskInstance: { region_id: regionId,
region_index: 3 } });
+ await act(async () => releasePass2());
+
+ assert.match(plugin.text(), /model-for-pass-3/);
+ assert.doesNotMatch(plugin.text(), /model-for-pass-2/);
+});
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/region.test.ts
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/region.test.ts
new file mode 100644
index 00000000000..9423061a149
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/region.test.ts
@@ -0,0 +1,45 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import assert from "node:assert/strict";
+import { test } from "node:test";
+
+import { isSelectionPending, regionOf, regionQuery } from "../src/region.ts";
+
+const regionId = "00000000-0000-0000-0000-000000000123";
+
+for (const regionIndex of [-1, 0, 2]) {
+ test(`a loop task instance at region index ${regionIndex} addresses its
pass`, () => {
+ const region = regionOf({ id: "ti", region_id: regionId, region_index:
regionIndex });
+
+ assert.deepEqual(region, { region_id: regionId, region_index: regionIndex
});
+ assert.equal(regionQuery(region),
`®ion_id=${regionId}®ion_index=${regionIndex}`);
+ });
+}
+
+test("a task instance outside a loop addresses no region", () => {
+ assert.equal(regionOf({ id: "ti" }), undefined);
+ assert.equal(regionQuery(regionOf({ id: "ti" })), "");
+});
+
+test("only a host that passes an unresolved task instance is pending", () => {
+ assert.equal(isSelectionPending({ taskInstance: undefined }), true);
+ assert.equal(isSelectionPending({ taskInstance: { id: "ti" } }), false);
+ assert.equal(isSelectionPending({}), false);
+});
