From c710d60b25f831957ee47e18781ce3c76686ed92 Mon Sep 17 00:00:00 2001
From: Donnie Adams <donnie@acorn.io>
Date: Wed, 23 Oct 2024 15:24:16 -0400
Subject: [PATCH] feat: detect not found errors in workspace API

This will allow for proper error handling in the workspace API.

Signed-off-by: Donnie Adams <donnie@acorn.io>
---
 workspace.go      | 35 ++++++++++++++++++++++++++++++++---
 workspace_test.go | 13 +++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/workspace.go b/workspace.go
index 6d79998..cc51e58 100644
--- a/workspace.go
+++ b/workspace.go
@@ -2,10 +2,26 @@ package gptscript
 
 import (
 	"context"
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
 	"os"
 	"strings"
 )
 
+type NotFoundInWorkspaceError struct {
+	id   string
+	name string
+}
+
+func (e *NotFoundInWorkspaceError) Error() string {
+	return fmt.Sprintf("not found: %s/%s", e.id, e.name)
+}
+
+func newNotFoundInWorkspaceError(id, name string) *NotFoundInWorkspaceError {
+	return &NotFoundInWorkspaceError{id: id, name: name}
+}
+
 func (g *GPTScript) CreateWorkspace(ctx context.Context, providerType string, fromWorkspaces ...string) (string, error) {
 	out, err := g.runBasicCommand(ctx, "workspaces/create", map[string]any{
 		"providerType":     providerType,
@@ -75,7 +91,13 @@ func (g *GPTScript) ListFilesInWorkspace(ctx context.Context, opts ...ListFilesI
 		return nil, err
 	}
 
-	return strings.Split(strings.TrimSpace(out), "\n"), nil
+	out = strings.TrimSpace(out)
+	if len(out) == 0 {
+		return nil, nil
+	}
+
+	var files []string
+	return files, json.Unmarshal([]byte(out), &files)
 }
 
 type RemoveAllOptions struct {
@@ -126,7 +148,7 @@ func (g *GPTScript) WriteFileInWorkspace(ctx context.Context, filePath string, c
 
 	_, err := g.runBasicCommand(ctx, "workspaces/write-file", map[string]any{
 		"id":            opt.WorkspaceID,
-		"contents":      contents,
+		"contents":      base64.StdEncoding.EncodeToString(contents),
 		"filePath":      filePath,
 		"workspaceTool": g.globalOpts.WorkspaceTool,
 		"env":           g.globalOpts.Env,
@@ -158,6 +180,10 @@ func (g *GPTScript) DeleteFileInWorkspace(ctx context.Context, filePath string,
 		"env":           g.globalOpts.Env,
 	})
 
+	if err != nil && strings.HasSuffix(err.Error(), fmt.Sprintf("not found: %s/%s", opt.WorkspaceID, filePath)) {
+		return newNotFoundInWorkspaceError(opt.WorkspaceID, filePath)
+	}
+
 	return err
 }
 
@@ -184,8 +210,11 @@ func (g *GPTScript) ReadFileInWorkspace(ctx context.Context, filePath string, op
 		"env":           g.globalOpts.Env,
 	})
 	if err != nil {
+		if strings.HasSuffix(err.Error(), fmt.Sprintf("not found: %s/%s", opt.WorkspaceID, filePath)) {
+			return nil, newNotFoundInWorkspaceError(opt.WorkspaceID, filePath)
+		}
 		return nil, err
 	}
 
-	return []byte(out), nil
+	return base64.StdEncoding.DecodeString(out)
 }
diff --git a/workspace_test.go b/workspace_test.go
index a7ff08c..8881d00 100644
--- a/workspace_test.go
+++ b/workspace_test.go
@@ -3,6 +3,7 @@ package gptscript
 import (
 	"bytes"
 	"context"
+	"errors"
 	"os"
 	"testing"
 )
@@ -46,6 +47,12 @@ func TestWriteReadAndDeleteFileFromWorkspace(t *testing.T) {
 		t.Errorf("Unexpected content: %s", content)
 	}
 
+	// Ensure we get the error we expect when trying to read a non-existent file
+	_, err = g.ReadFileInWorkspace(context.Background(), "test1.txt", ReadFileInWorkspaceOptions{WorkspaceID: id})
+	if nf := (*NotFoundInWorkspaceError)(nil); !errors.As(err, &nf) {
+		t.Errorf("Unexpected error reading non-existent file: %v", err)
+	}
+
 	err = g.DeleteFileInWorkspace(context.Background(), "test.txt", DeleteFileInWorkspaceOptions{WorkspaceID: id})
 	if err != nil {
 		t.Errorf("Error deleting file: %v", err)
@@ -169,6 +176,12 @@ func TestWriteReadAndDeleteFileFromWorkspaceS3(t *testing.T) {
 		t.Errorf("Unexpected content: %s", content)
 	}
 
+	// Ensure we get the error we expect when trying to read a non-existent file
+	_, err = g.ReadFileInWorkspace(context.Background(), "test1.txt", ReadFileInWorkspaceOptions{WorkspaceID: id})
+	if nf := (*NotFoundInWorkspaceError)(nil); !errors.As(err, &nf) {
+		t.Errorf("Unexpected error reading non-existent file: %v", err)
+	}
+
 	err = g.DeleteFileInWorkspace(context.Background(), "test.txt", DeleteFileInWorkspaceOptions{WorkspaceID: id})
 	if err != nil {
 		t.Errorf("Error deleting file: %v", err)