From 028fbcfa68f6340286f62f0cb3b61d99a701cc5f Mon Sep 17 00:00:00 2001
From: Michael Doyle <michaeldoyle@google.com>
Date: Tue, 24 Feb 2026 17:17:35 -0500
Subject: [PATCH] feat(js): add longRunning support and refine model types

---
 js/ai/src/model-types.ts                    |  2 ++
 js/ai/src/model.ts                          | 12 +++++++++++
 js/ai/src/parts.ts                          |  9 +++++++--
 js/plugins/google-genai/src/googleai/veo.ts | 22 ++++++++++-----------
 js/plugins/google-genai/src/vertexai/veo.ts | 22 ++++++++++-----------
 5 files changed, 41 insertions(+), 26 deletions(-)

diff --git a/js/ai/src/model-types.ts b/js/ai/src/model-types.ts
index d76a555c9a..1bf3da805b 100644
--- a/js/ai/src/model-types.ts
+++ b/js/ai/src/model-types.ts
@@ -106,6 +106,8 @@ export const ModelInfoSchema = z.object({
       constrained: z.enum(['none', 'all', 'no-tools']).optional(),
       /** Model supports controlling tool choice, e.g. forced tool calling. */
       toolChoice: z.boolean().optional(),
+      /** Model can perform long-running operations. */
+      longRunning: z.boolean().optional(),
     })
     .optional(),
   /** At which stage of development this model is.
diff --git a/js/ai/src/model.ts b/js/ai/src/model.ts
index 442576e3d8..5ed77f3b8e 100644
--- a/js/ai/src/model.ts
+++ b/js/ai/src/model.ts
@@ -57,12 +57,18 @@ import {
   CustomPartSchema,
   DataPartSchema,
   MediaPartSchema,
+  MultipartToolResponseSchema,
+  ReasoningPartSchema,
+  ResourcePartSchema,
   TextPartSchema,
   ToolRequestPartSchema,
   ToolResponsePartSchema,
   type CustomPart,
   type DataPart,
   type MediaPart,
+  type MultipartToolResponse,
+  type ReasoningPart,
+  type ResourcePart,
   type TextPart,
   type ToolRequestPart,
   type ToolResponsePart,
@@ -73,6 +79,9 @@ export {
   CustomPartSchema,
   DataPartSchema,
   MediaPartSchema,
+  MultipartToolResponseSchema,
+  ReasoningPartSchema,
+  ResourcePartSchema,
   TextPartSchema,
   ToolRequestPartSchema,
   ToolResponsePartSchema,
@@ -80,6 +89,9 @@ export {
   type CustomPart,
   type DataPart,
   type MediaPart,
+  type MultipartToolResponse,
+  type ReasoningPart,
+  type ResourcePart,
   type TextPart,
   type ToolRequestPart,
   type ToolResponsePart,
diff --git a/js/ai/src/parts.ts b/js/ai/src/parts.ts
index 3a81cd9660..53fe2f086f 100644
--- a/js/ai/src/parts.ts
+++ b/js/ai/src/parts.ts
@@ -36,6 +36,11 @@ export const TextPartSchema = EmptyPartSchema.extend({
   text: z.string(),
 });
 
+/**
+ * Text part.
+ */
+export type TextPart = z.infer<typeof TextPartSchema>;
+
 /**
  * Zod schema for a reasoning part.
  */
@@ -45,9 +50,9 @@ export const ReasoningPartSchema = EmptyPartSchema.extend({
 });
 
 /**
- * Text part.
+ * Reasoning part.
  */
-export type TextPart = z.infer<typeof TextPartSchema>;
+export type ReasoningPart = z.infer<typeof ReasoningPartSchema>;
 
 /**
  * Zod schema of media.
diff --git a/js/plugins/google-genai/src/googleai/veo.ts b/js/plugins/google-genai/src/googleai/veo.ts
index 80f81cb3cc..e563541ab8 100644
--- a/js/plugins/google-genai/src/googleai/veo.ts
+++ b/js/plugins/google-genai/src/googleai/veo.ts
@@ -102,18 +102,16 @@ function commonRef(
   return modelRef({
     name: `googleai/${name}`,
     configSchema,
-    info:
-      info ??
-      ({
-        supports: {
-          media: true,
-          multiturn: false,
-          tools: false,
-          systemRole: false,
-          output: ['media'],
-          longRunning: true,
-        },
-      } as ModelInfo), // TODO(ifielker): Remove this cast if we fix longRunning
+    info: info ?? {
+      supports: {
+        media: true,
+        multiturn: false,
+        tools: false,
+        systemRole: false,
+        output: ['media'],
+        longRunning: true,
+      },
+    },
   });
 }
 
diff --git a/js/plugins/google-genai/src/vertexai/veo.ts b/js/plugins/google-genai/src/vertexai/veo.ts
index 969da00c69..c3a7537491 100644
--- a/js/plugins/google-genai/src/vertexai/veo.ts
+++ b/js/plugins/google-genai/src/vertexai/veo.ts
@@ -129,18 +129,16 @@ function commonRef(
   return modelRef({
     name: `vertexai/${name}`,
     configSchema,
-    info:
-      info ??
-      ({
-        supports: {
-          media: true,
-          multiturn: false,
-          tools: false,
-          systemRole: false,
-          output: ['media'],
-          longRunning: true,
-        },
-      } as ModelInfo), // TODO(ifielker): Remove this cast if we fix longRunning
+    info: info ?? {
+      supports: {
+        media: true,
+        multiturn: false,
+        tools: false,
+        systemRole: false,
+        output: ['media'],
+        longRunning: true,
+      },
+    },
   });
 }