refactor: tighten internal types and eslint fixes (assistant-ui#1911)

Author: Yonom

packages/assistant-stream/src/core/tool/ToolCallReader.ts

@@ -8,18 +8,24 @@ import {
   ToolCallReader,
   ToolCallResponseReader,
 } from "./tool-types";
-import { TypeAtPath, TypePath } from "./type-path-utils";
+import { DeepPartial, TypeAtPath, TypePath } from "./type-path-utils";
 import { ToolResponse } from "./ToolResponse";
+import {
+  asAsyncIterableStream,
+  AsyncIterableStream,
+  ReadonlyJSONObject,
+  ReadonlyJSONValue,
+} from "../../utils";
 
 // TODO: remove dispose
 
-function getField<T>(obj: T, fieldPath: (string | number)[]): any {
-  let current: any = obj;
+function getField<T>(obj: T, fieldPath: (string | number)[]): unknown {
+  let current: unknown = obj;
   for (const key of fieldPath) {
     if (current === undefined || current === null) {
       return undefined;
     }
-    current = current[key as string | number];
+    current = current[key as keyof typeof current];
   }
   return current;
 }
@@ -29,14 +35,14 @@ interface Handle {
   dispose(): void;
 }
 
-class GetHandle<T> implements Handle {
-  private resolve: (value: any) => void;
+class GetHandle<T, TValue> implements Handle {
+  private resolve: (value: TValue) => void;
   private reject: (reason: unknown) => void;
   private disposed = false;
   private fieldPath: (string | number)[];
 
   constructor(
-    resolve: (value: any) => void,
+    resolve: (value: TValue) => void,
     reject: (reason: unknown) => void,
     fieldPath: (string | number)[],
   ) {
@@ -58,7 +64,7 @@ class GetHandle<T> implements Handle {
       ) {
         const value = getField(args as T, this.fieldPath);
         if (value !== undefined) {
-          this.resolve(value);
+          this.resolve(value as TValue);
           this.dispose();
         }
       }
@@ -74,12 +80,12 @@ class GetHandle<T> implements Handle {
 }
 
 class StreamValuesHandle<T> implements Handle {
-  private controller: ReadableStreamDefaultController<any>;
+  private controller: ReadableStreamDefaultController<unknown>;
   private disposed = false;
   private fieldPath: (string | number)[];
 
   constructor(
-    controller: ReadableStreamDefaultController<any>,
+    controller: ReadableStreamDefaultController<unknown>,
     fieldPath: (string | number)[],
   ) {
     this.controller = controller;
@@ -118,13 +124,13 @@ class StreamValuesHandle<T> implements Handle {
 }
 
 class StreamTextHandle<T> implements Handle {
-  private controller: ReadableStreamDefaultController<any>;
+  private controller: ReadableStreamDefaultController<unknown>;
   private disposed = false;
   private fieldPath: (string | number)[];
-  private lastValue: any = undefined;
+  private lastValue: string | undefined = undefined;
 
   constructor(
-    controller: ReadableStreamDefaultController<any>,
+    controller: ReadableStreamDefaultController<unknown>,
     fieldPath: (string | number)[],
   ) {
     this.controller = controller;
@@ -165,13 +171,13 @@ class StreamTextHandle<T> implements Handle {
 }
 
 class ForEachHandle<T> implements Handle {
-  private controller: ReadableStreamDefaultController<any>;
+  private controller: ReadableStreamDefaultController<unknown>;
   private disposed = false;
   private fieldPath: (string | number)[];
   private processedIndexes = new Set<number>();
 
   constructor(
-    controller: ReadableStreamDefaultController<any>,
+    controller: ReadableStreamDefaultController<unknown>,
     fieldPath: (string | number)[],
   ) {
     this.controller = controller;
@@ -182,7 +188,7 @@ class ForEachHandle<T> implements Handle {
     if (this.disposed) return;
 
     try {
-      const array = getField(args as T, this.fieldPath) as unknown as any[];
+      const array = getField(args as T, this.fieldPath);
 
       if (!Array.isArray(array)) {
         return;
@@ -226,10 +232,12 @@ class ForEachHandle<T> implements Handle {
 }
 
 // Implementation of ToolCallReader that uses stream of partial JSON
-export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
+export class ToolCallArgsReaderImpl<T extends ReadonlyJSONObject>
+  implements ToolCallArgsReader<T>
+{
   private argTextDeltas: ReadableStream<string>;
   private handles: Set<Handle> = new Set();
-  private args: any = parsePartialJsonObject("");
+  private args: unknown = parsePartialJsonObject("");
 
   constructor(argTextDeltas: ReadableStream<string>) {
     this.argTextDeltas = argTextDeltas;
@@ -268,8 +276,12 @@ export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
   get<PathT extends TypePath<T>>(
     ...fieldPath: PathT
   ): Promise<TypeAtPath<T, PathT>> {
-    return new Promise<any>((resolve, reject) => {
-      const handle = new GetHandle<T>(resolve, reject, fieldPath);
+    return new Promise<TypeAtPath<T, PathT>>((resolve, reject) => {
+      const handle = new GetHandle<T, TypeAtPath<T, PathT>>(
+        resolve,
+        reject,
+        fieldPath,
+      );
 
       // Check if the field is already complete in current args
       if (
@@ -281,7 +293,7 @@ export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
       ) {
         const value = getField(this.args as T, fieldPath);
         if (value !== undefined) {
-          resolve(value);
+          resolve(value as TypeAtPath<T, PathT>);
           return;
         }
       }
@@ -291,11 +303,13 @@ export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
     });
   }
 
-  streamValues<PathT extends TypePath<T>>(...fieldPath: PathT): any {
+  streamValues<PathT extends TypePath<T>>(
+    ...fieldPath: PathT
+  ): AsyncIterableStream<DeepPartial<TypeAtPath<T, PathT>>> {
     // Use a type assertion to convert the complex TypePath to a simple array
     const simplePath = fieldPath as unknown as (string | number)[];
 
-    const stream = new ReadableStream<any>({
+    const stream = new ReadableStream<DeepPartial<TypeAtPath<T, PathT>>>({
       start: (controller) => {
         const handle = new StreamValuesHandle<T>(controller, simplePath);
         this.handles.add(handle);
@@ -315,15 +329,19 @@ export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
       },
     });
 
-    // For type compatibility, cast the stream to the required type
-    return stream as any;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    return asAsyncIterableStream(stream) as any;
   }
 
-  streamText<PathT extends TypePath<T>>(...fieldPath: PathT): any {
+  streamText<PathT extends TypePath<T>>(
+    ...fieldPath: PathT
+  ): TypeAtPath<T, PathT> extends string & infer U
+    ? AsyncIterableStream<U>
+    : never {
     // Use a type assertion to convert the complex TypePath to a simple array
     const simplePath = fieldPath as unknown as (string | number)[];
 
-    const stream = new ReadableStream<any>({
+    const stream = new ReadableStream<unknown>({
       start: (controller) => {
         const handle = new StreamTextHandle<T>(controller, simplePath);
         this.handles.add(handle);
@@ -343,15 +361,19 @@ export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
       },
     });
 
-    // For type compatibility, cast the stream to the required type
-    return stream as any;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    return asAsyncIterableStream(stream) as any;
   }
 
-  forEach<PathT extends TypePath<T>>(...fieldPath: PathT): any {
+  forEach<PathT extends TypePath<T>>(
+    ...fieldPath: PathT
+  ): TypeAtPath<T, PathT> extends Array<infer U>
+    ? AsyncIterableStream<U>
+    : never {
     // Use a type assertion to convert the complex TypePath to a simple array
     const simplePath = fieldPath as unknown as (string | number)[];
 
-    const stream = new ReadableStream<any>({
+    const stream = new ReadableStream<unknown>({
       start: (controller) => {
         const handle = new ForEachHandle<T>(controller, simplePath);
         this.handles.add(handle);
@@ -371,12 +393,12 @@ export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
       },
     });
 
-    // For type compatibility, cast the stream to the required type
-    return stream as any;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    return asAsyncIterableStream(stream) as any;
   }
 }
 
-export class ToolCallResponseReaderImpl<TResult>
+export class ToolCallResponseReaderImpl<TResult extends ReadonlyJSONValue>
   implements ToolCallResponseReader<TResult>
 {
   constructor(private readonly promise: Promise<ToolResponse<TResult>>) {}
@@ -386,8 +408,10 @@ export class ToolCallResponseReaderImpl<TResult>
   }
 }
 
-export class ToolCallReaderImpl<TArgs, TResult>
-  implements ToolCallReader<TArgs, TResult>
+export class ToolCallReaderImpl<
+  TArgs extends ReadonlyJSONObject,
+  TResult extends ReadonlyJSONValue,
+> implements ToolCallReader<TArgs, TResult>
 {
   public readonly args: ToolCallArgsReaderImpl<TArgs>;
   public readonly response: ToolCallResponseReaderImpl<TResult>;

packages/assistant-stream/src/core/tool/ToolExecutionStream.ts

@@ -5,7 +5,10 @@ import {
   AssistantMetaTransformStream,
 } from "../utils/stream/AssistantMetaTransformStream";
 import { PipeableTransformStream } from "../utils/stream/PipeableTransformStream";
-import { ReadonlyJSONValue } from "../../utils/json/json-value";
+import {
+  ReadonlyJSONObject,
+  ReadonlyJSONValue,
+} from "../../utils/json/json-value";
 import { ToolResponse } from "./ToolResponse";
 import { withPromiseOrValue } from "../utils/withPromiseOrValue";
 import { ToolCallReaderImpl } from "./ToolCallReader";
@@ -14,13 +17,16 @@ import { ToolCallReader } from "./tool-types";
 type ToolCallback = (toolCall: {
   toolCallId: string;
   toolName: string;
-  args: unknown;
+  args: ReadonlyJSONObject;
 }) =>
   | Promise<ToolResponse<ReadonlyJSONValue>>
   | ToolResponse<ReadonlyJSONValue>
   | undefined;
 
-type ToolStreamCallback = <TArgs, TResult>(toolCall: {
+type ToolStreamCallback = <
+  TArgs extends ReadonlyJSONObject = ReadonlyJSONObject,
+  TResult extends ReadonlyJSONValue = ReadonlyJSONValue,
+>(toolCall: {
   reader: ToolCallReader<TArgs, TResult>;
   toolCallId: string;
   toolName: string;
@@ -39,7 +45,7 @@ export class ToolExecutionStream extends PipeableTransformStream<
     const toolCallPromises = new Map<string, PromiseLike<void>>();
     const toolCallControllers = new Map<
       string,
-      ToolCallReaderImpl<unknown, unknown>
+      ToolCallReaderImpl<ReadonlyJSONObject, ReadonlyJSONValue>
     >();
 
     super((readable) => {
@@ -58,7 +64,10 @@ export class ToolExecutionStream extends PipeableTransformStream<
           switch (type) {
             case "part-start":
               if (chunk.part.type === "tool-call") {
-                const reader = new ToolCallReaderImpl<unknown, unknown>();
+                const reader = new ToolCallReaderImpl<
+                  ReadonlyJSONObject,
+                  ReadonlyJSONValue
+                >();
                 toolCallControllers.set(chunk.part.toolCallId, reader);
 
                 options.streamCall({

packages/assistant-stream/src/core/tool/ToolResponse.ts

@@ -23,7 +23,9 @@ export class ToolResponse<TResult> {
     this.isError = options.isError ?? false;
   }
 
-  static [Symbol.hasInstance](obj: unknown): obj is ToolResponse<unknown> {
+  static [Symbol.hasInstance](
+    obj: unknown,
+  ): obj is ToolResponse<ReadonlyJSONValue> {
     return (
       typeof obj === "object" && obj !== null && TOOL_RESPONSE_SYMBOL in obj
     );

packages/assistant-stream/src/core/tool/tool-types.ts

@@ -1,6 +1,5 @@
 import { JSONSchema7 } from "json-schema";
-import { TypeAtPath, TypePath } from "./type-path-utils";
-import { DeepPartial } from "ai";
+import { DeepPartial, TypeAtPath, TypePath } from "./type-path-utils";
 import { AsyncIterableStream } from "../../utils";
 import type { StandardSchemaV1 } from "@standard-schema/spec";
 import { ToolResponse } from "./ToolResponse";
@@ -13,7 +12,7 @@ import { ToolResponse } from "./ToolResponse";
  *
  * @template TArgs The type of arguments being read.
  */
-export interface ToolCallArgsReader<TArgs> {
+export interface ToolCallArgsReader<TArgs extends Record<string, unknown>> {
   /**
    * Returns a promise that will resolve to the value at the given path,
    * as soon as that path is generated by the LLM.
@@ -63,7 +62,10 @@ export interface ToolCallResponseReader<TResult> {
   get: () => Promise<ToolResponse<TResult>>;
 }
 
-export interface ToolCallReader<TArgs, TResult> {
+export interface ToolCallReader<
+  TArgs extends Record<string, unknown> = Record<string, unknown>,
+  TResult = unknown,
+> {
   args: ToolCallArgsReader<TArgs>;
   response: ToolCallResponseReader<TResult>;
 
@@ -85,7 +87,10 @@ export type ToolExecuteFunction<TArgs, TResult> = (
   context: ToolExecutionContext,
 ) => TResult | Promise<TResult>;
 
-export type ToolStreamCallFunction<TArgs, TResult> = (
+export type ToolStreamCallFunction<
+  TArgs extends Record<string, unknown> = Record<string, unknown>,
+  TResult = unknown,
+> = (
   reader: ToolCallReader<TArgs, TResult>,
   context: ToolExecutionContext,
 ) => void;