fixes

Author: maticzav

examples/demo.ts

@@ -1,7 +1,6 @@
 #!/usr/bin/env -S npm run tsn -T
 
 import { BrowserUse } from 'browser-use-sdk';
-import { TaskView } from 'browser-use-sdk/resources';
 import { spinner } from './utils';
 
 // gets API Key from environment variable BROWSER_USE_API_KEY
@@ -18,7 +17,7 @@ async function main() {
 
   poll: do {
     // Wait for Task to Finish
-    const status = (await browseruse.tasks.retrieve(rsp.id, { statusOnly: false })) as TaskView;
+    const status = await browseruse.tasks.retrieve(rsp.id);
 
     switch (status.status) {
       case 'started':

examples/structured-output.ts

@@ -23,15 +23,16 @@ async function main() {
   const stop = spinner(() => log);
 
   // Create Task
-  const rsp = await browseruse.tasks.createWithStructuredOutput({
+  const rsp = await browseruse.tasks.create({
     task: 'Extract top 10 Hacker News posts and return the title, url, and score',
     structuredOutputJson: TaskOutput,
   });
 
   poll: do {
     // Wait for Task to Finish
-    const status = await browseruse.tasks.retrieveWithStructuredOutput(rsp.id, {
-      structuredOutputJson: TaskOutput,
+    const status = await browseruse.tasks.retrieve({
+      taskId: rsp.id,
+      schema: TaskOutput,
     });
 
     switch (status.status) {
@@ -52,12 +53,16 @@ async function main() {
       }
 
       case 'finished':
+        if (status.doneOutput == null) {
+          throw new Error('No output');
+        }
+
         stop();
 
         // Print Structured Output
         console.log('TOP POSTS:');
 
-        for (const post of status.doneOutput!.posts) {
+        for (const post of status.doneOutput.posts) {
           console.log(` - ${post.title} (${post.score}) ${post.url}`);
         }
 

src/lib/parse.ts

@@ -1,46 +1,34 @@
 import z, { type ZodType } from 'zod';
-import type { TaskCreateParams, TaskRetrieveParams, TaskView } from '../resources/tasks';
+import type { TaskCreateParams, TaskView } from '../resources/tasks';
 
 // RUN
 
-export type RunTaskCreateParamsWithStructuredOutput<T extends ZodType> = Omit<
-  TaskCreateParams,
-  'structuredOutputJson'
-> & {
+export type TaskCreateParamsWithSchema<T extends ZodType> = Omit<TaskCreateParams, 'structuredOutputJson'> & {
   structuredOutputJson: T;
 };
 
-export function stringifyStructuredOutput<T extends ZodType>(
-  req: RunTaskCreateParamsWithStructuredOutput<T>,
-): TaskCreateParams {
-  return {
-    ...req,
-    structuredOutputJson: JSON.stringify(z.toJSONSchema(req.structuredOutputJson)),
-  };
+export function stringifyStructuredOutput<T extends ZodType>(schema: T): string {
+  return JSON.stringify(z.toJSONSchema(schema));
 }
 
 // RETRIEVE
 
-export type GetTaskStatusParamsWithStructuredOutput<T extends ZodType> = Omit<
-  TaskRetrieveParams,
-  'statusOnly'
-> & {
-  statusOnly?: false;
-  structuredOutputJson: T;
-};
-
-export type TaskViewWithStructuredOutput<T extends ZodType> = Omit<TaskView, 'doneOutput'> & {
+export type TaskViewWithSchema<T extends ZodType> = Omit<TaskView, 'doneOutput'> & {
   doneOutput: z.output<T> | null;
 };
 
 export function parseStructuredTaskOutput<T extends ZodType>(
   res: TaskView,
-  body: GetTaskStatusParamsWithStructuredOutput<T>,
-): TaskViewWithStructuredOutput<T> {
+  schema: T,
+): TaskViewWithSchema<T> {
+  if (res.doneOutput == null) {
+    return { ...res, doneOutput: null };
+  }
+
   try {
     const parsed = JSON.parse(res.doneOutput);
 
-    const response = body.structuredOutputJson.safeParse(parsed);
+    const response = schema.safeParse(parsed);
     if (!response.success) {
       throw new Error(`Invalid structured output: ${response.error.message}`);
     }

src/resources/tasks.ts

@@ -10,9 +10,8 @@ import { path } from '../internal/utils/path';
 import {
   parseStructuredTaskOutput,
   stringifyStructuredOutput,
-  type TaskViewWithStructuredOutput,
-  type GetTaskStatusParamsWithStructuredOutput,
-  type RunTaskCreateParamsWithStructuredOutput,
+  type TaskViewWithSchema,
+  type TaskCreateParamsWithSchema,
 } from '../lib/parse';
 
 export class Tasks extends APIResource {
@@ -62,17 +61,33 @@ export class Tasks extends APIResource {
    * - 404: If referenced agent/browser profiles don't exist
    * - 400: If session is stopped or already has a running task
    */
-  create(body: TaskCreateParams, options?: RequestOptions): APIPromise<TaskView> {
-    return this._client.post('/tasks', { body, ...options });
-  }
-
-  createWithStructuredOutput<T extends ZodType>(
-    body: RunTaskCreateParamsWithStructuredOutput<T>,
+  create<T extends ZodType>(
+    body: TaskCreateParamsWithSchema<T>,
     options?: RequestOptions,
-  ): APIPromise<TaskViewWithStructuredOutput<T>> {
-    return this.create(stringifyStructuredOutput(body), options)._thenUnwrap((rsp) =>
-      parseStructuredTaskOutput(rsp as TaskView, body),
-    );
+  ): APIPromise<TaskViewWithSchema<T>>;
+  create(body: TaskCreateParams, options?: RequestOptions): APIPromise<TaskView>;
+  create(
+    body: TaskCreateParams | TaskCreateParamsWithSchema<ZodType>,
+    options?: RequestOptions,
+  ): APIPromise<unknown> {
+    if (body.structuredOutputJson == null || typeof body.structuredOutputJson === 'string') {
+      return this._client.post('/tasks', { body, ...options });
+    }
+
+    if (typeof body.structuredOutputJson === 'object') {
+      const schema = body.structuredOutputJson;
+
+      const _body: TaskCreateParams = {
+        ...body,
+        structuredOutputJson: stringifyStructuredOutput(schema),
+      };
+
+      return this._client
+        .post('/tasks', { body: _body, ...options })
+        ._thenUnwrap((rsp) => parseStructuredTaskOutput(rsp as TaskView, schema));
+    }
+
+    return this._client.post('/tasks', { body, ...options });
   }
 
   /**
@@ -101,22 +116,21 @@ export class Tasks extends APIResource {
    *
    * - 404: If the user agent task doesn't exist
    */
-  retrieve(taskID: string, options?: RequestOptions): APIPromise<TaskView> {
-    return this._client.get(path`/tasks/${taskID}`, options);
-  }
-
-  retrieveWithStructuredOutput<T extends ZodType>(
-    taskID: string,
-    query: GetTaskStatusParamsWithStructuredOutput<T>,
+  retrieve<T extends ZodType>(
+    req: { taskId: string; schema: T },
     options?: RequestOptions,
-  ): APIPromise<TaskViewWithStructuredOutput<T>> {
-    // NOTE: We manually remove structuredOutputJson from the query object because
-    //       it's not a valid Browser Use Cloud parameter.
-    const { structuredOutputJson, ...rest } = query;
-
-    return this.retrieve(taskID, rest, options)._thenUnwrap((rsp) =>
-      parseStructuredTaskOutput(rsp as TaskView, query),
-    );
+  ): APIPromise<TaskViewWithSchema<T>>;
+  retrieve(taskID: string, options?: RequestOptions): APIPromise<TaskView>;
+  retrieve(req: string | { taskId: string; schema: ZodType }, options?: RequestOptions): APIPromise<unknown> {
+    if (typeof req === 'string') {
+      return this._client.get(path`/tasks/${req}`, options);
+    }
+
+    const { taskId, schema } = req;
+
+    return this._client
+      .get(path`/tasks/${taskId}`, options)
+      ._thenUnwrap((rsp) => parseStructuredTaskOutput(rsp as TaskView, schema));
   }
 
   /**