dagger · vito · Jan 7, 2025 · Jan 7, 2025 · Jan 7, 2025 · Jan 7, 2025
diff --git a/.dagger/main.go b/.dagger/main.go
@@ -66,7 +66,7 @@ func New(
 		if strings.HasPrefix(module, "core/integration/") {
 			continue
 		}
-		if strings.HasPrefix(module, "dagql/idtui/viztest/broken") {
+		if strings.HasPrefix(module, "dagql/idtui/viztest/") {
 			continue
 		}
 		if strings.HasPrefix(module, "modules/evaluator/") {
@@ -120,7 +120,7 @@ func (dev *DaggerDev) Lint(
 				if strings.HasPrefix(pkg, "core/integration/") {
 					continue
 				}
-				if strings.HasPrefix(pkg, "dagql/idtui/viztest/broken") {
+				if strings.HasPrefix(pkg, "dagql/idtui/viztest/") {
 					continue
 				}
 				if strings.HasPrefix(pkg, "modules/evaluator/") {

diff --git a/cmd/codegen/generator/go/templates/src/_types/object.go.tmpl b/cmd/codegen/generator/go/templates/src/_types/object.go.tmpl
@@ -3,11 +3,12 @@
 type {{ .Name | FormatName }} struct {
 {{- with .Directives.SourceMap -}} // {{ .Module }} ({{ .Filelink | ModuleRelPath }}) {{- end }}
 	query *querybuilder.Selection
+	client graphql.Client
 
-    {{ range $field := .Fields }}
-        {{- if $field.TypeRef.IsScalar }}
-        {{ $field.Name }} *{{ $field.TypeRef | FormatOutputType }}
-        {{- end }}
+	{{ range $field := .Fields }}
+	{{- if $field.TypeRef.IsScalar }}
+	{{ $field.Name }} *{{ $field.TypeRef | FormatOutputType }}
+	{{- end }}
 	{{- end }}
 }
 {{- end }}
@@ -28,12 +29,64 @@ func (r *{{ $.Name | FormatName }}) With(f With{{ .Name | FormatName }}Func) *{{
 func (r *{{ .Name | FormatName }}) WithGraphQLQuery(q *querybuilder.Selection) *{{ $.Name | FormatName }} {
 	return &{{ .Name | FormatName }}{
 		query: q,
-		{{- if eq .Name "Query" }}
 		client: r.client,
-		{{ end }}
 	}
 }
 
+{{- if eq .Name "Status" }}
+func (r *Status) Context(ctx context.Context) (context.Context, *Status) {
+	started, err := r.Start(ctx)
+	if err != nil {
+		panic(err)
+	}
+	spanIDHex, err := started.InternalID(ctx)
+	if err != nil {
+		panic(err)
+	}
+	spanID, err := trace.SpanIDFromHex(spanIDHex)
+	if err != nil {
+		panic(err)
+	}
+	spanCtx := trace.SpanContextFromContext(ctx)
+	return trace.ContextWithSpanContext(ctx, trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID:    spanCtx.TraceID(),
+		SpanID:     spanID,
+		Remote:     true,
+		TraceFlags: spanCtx.TraceFlags(),
+		TraceState: spanCtx.TraceState(),
+	})), started
+}
+
+{{- $supportsVoid := CheckVersionCompatibility "v0.12.0" }}
+func (r *Status) Run(ctx context.Context, cb func(context.Context) error) error {
+	ctx, status := r.Context(ctx)
+	err := cb(ctx)
+	var endErr error
+	if err != nil {
+		errClient := &Client{
+			query: querybuilder.Query().Client(r.client),
+			client: r.client,
+		}
+		{{- if $supportsVoid }}
+		endErr = status.End(ctx, StatusEndOpts{
+			Error: errClient.Error(err.Error()),
+		})
+		{{- else }}
+		_, endErr = status.End(ctx, StatusEndOpts{
+			Error: errClient.Error(err.Error()),
+		})
+		{{- end }}
+	} else {
+		{{- if $supportsVoid }}
+		endErr = status.End(ctx)
+		{{- else }}
+		_, endErr = status.End(ctx)
+		{{- end }}
+	}
+	return errors.Join(err, endErr)
+}
+{{- end }}
+
 {{ range $field := .Fields }}
 {{- if HasOptionals $field.Args }}
 // {{ $field | FieldOptionsStructName }} contains options for {{ $.Name | FormatName }}.{{ $field.Name | FormatName }}
@@ -75,32 +128,32 @@ type {{ $field | FieldOptionsStructName }} struct {
 {{ FieldFunction $field false $supportsVoid }} {
 {{- with .Directives.SourceMap -}} // {{ .Module }} ({{ .Filelink | ModuleRelPath }}) {{- end }}
 	{{- range $arg := $field.Args }}
-	    {{- if and (IsPointer $arg) (not (IsArgOptional $arg)) }}
-        assertNotNil("{{ $arg.Name}}", {{ $arg.Name }})
-        {{- end }}
-    {{- end }}
-
-    {{- if and ($field.TypeRef.IsScalar) (ne $field.ParentObject.Name "Query") (not $convertID) }}
-    if r.{{ $field.Name }} != nil {
-        {{- if and $supportsVoid $field.TypeRef.IsVoid }}
-        return nil
-        {{- else }}
-        return *r.{{ $field.Name }}, nil
-        {{- end }}
-    }
-    {{- end }}
+	{{- if and (IsPointer $arg) (not (IsArgOptional $arg)) }}
+	assertNotNil("{{ $arg.Name}}", {{ $arg.Name }})
+	{{- end }}
+	{{- end }}
+
+	{{- if and ($field.TypeRef.IsScalar) (ne $field.ParentObject.Name "Query") (not $convertID) }}
+	if r.{{ $field.Name }} != nil {
+		{{- if and $supportsVoid $field.TypeRef.IsVoid }}
+		return nil
+		{{- else }}
+		return *r.{{ $field.Name }}, nil
+		{{- end }}
+	}
+	{{- end }}
 	q := r.query.Select("{{ $field.Name }}")
 
 	{{- if HasOptionals $field.Args }}
 	for i := len(opts) - 1; i >= 0; i-- {
-	{{- range $arg := $field.Args }}
-	{{- if IsArgOptional $arg }}
-	// `{{ $arg.Name }}` optional argument
-	if !querybuilder.IsZeroValue(opts[i].{{ $arg.Name | FormatName }}) {
-		q = q.Arg("{{ $arg.Name }}", opts[i].{{ $arg.Name | FormatName }})
-	}
-	{{- end }}
-	{{- end }}
+		{{- range $arg := $field.Args }}
+		{{- if IsArgOptional $arg }}
+		// `{{ $arg.Name }}` optional argument
+		if !querybuilder.IsZeroValue(opts[i].{{ $arg.Name | FormatName }}) {
+			q = q.Arg("{{ $arg.Name }}", opts[i].{{ $arg.Name | FormatName }})
+		}
+		{{- end }}
+		{{- end }}
 	}
 	{{- end }}
 
@@ -112,7 +165,7 @@ type {{ $field | FieldOptionsStructName }} struct {
 	{{- end }}
 	{{- $typeName := $field.TypeRef | FormatOutputType }}
 	{{ if and $supportsVoid $field.TypeRef.IsVoid }}
-		return q.Execute(ctx)
+	return q.Execute(ctx)
 	{{- else if $convertID }}
 	var id {{ $typeName }}
 	if err := q.Bind(&id).Execute(ctx); err != nil {
@@ -125,57 +178,55 @@ type {{ $field | FieldOptionsStructName }} struct {
 	{{- else if $field.TypeRef.IsObject }}
 	return &{{ $typeName }} {
 		query: q,
-		{{- if eq $typeName "Client" }}
 		client: r.client,
-		{{ end }}
 	}
 
 	{{- else if or $field.TypeRef.IsScalar $field.TypeRef.IsList }}
-		{{- if and $field.TypeRef.IsList (IsListOfObject $field.TypeRef) }}
-    q = q.Select("{{ range $i, $v := $field | GetArrayField }}{{ if $i }} {{ end }}{{ $v.Name }}{{ end }}")
-
-    type {{ $field.Name | ToLowerCase }} struct {
-      {{ range $v := $field | GetArrayField }}
-      {{ $v.Name | ToUpperCase }} {{ $v.TypeRef | FormatOutputType }}
-      {{- end }}
-    }
-
-    {{$eleType := $field.TypeRef | InnerType}}
-    convert := func(fields []{{ $field.Name | ToLowerCase }}) {{ $field.TypeRef | FormatOutputType }} {
-        out := {{ $field.TypeRef | FormatOutputType }}{}
-
-        for i := range fields {
-            val := {{ $field.TypeRef | FormatOutputType | FormatArrayToSingleType }}{{"{"}}{{ $field | GetArrayField | FormatArrayField }}{{"}"}}
-            {{- if $eleType | IsIDableObject }}
-              val.query = q.Root().Select("load{{$eleType | ObjectName}}FromID").Arg("id", fields[i].Id)
-            {{- end }}
-            out = append(out, val)
-        }
-
-        return out
-    }
-    {{- end }}
-
-    {{- if and $field.TypeRef.IsList (IsListOfObject $field.TypeRef) }}
+	{{- if and $field.TypeRef.IsList (IsListOfObject $field.TypeRef) }}
+	q = q.Select("{{ range $i, $v := $field | GetArrayField }}{{ if $i }} {{ end }}{{ $v.Name }}{{ end }}")
+
+	type {{ $field.Name | ToLowerCase }} struct {
+		{{ range $v := $field | GetArrayField }}
+		{{ $v.Name | ToUpperCase }} {{ $v.TypeRef | FormatOutputType }}
+		{{- end }}
+	}
+
+	{{$eleType := $field.TypeRef | InnerType}}
+	convert := func(fields []{{ $field.Name | ToLowerCase }}) {{ $field.TypeRef | FormatOutputType }} {
+		out := {{ $field.TypeRef | FormatOutputType }}{}
+
+		for i := range fields {
+			val := {{ $field.TypeRef | FormatOutputType | FormatArrayToSingleType }}{{"{"}}{{ $field | GetArrayField | FormatArrayField }}{{"}"}}
+			{{- if $eleType | IsIDableObject }}
+			val.query = q.Root().Select("load{{$eleType | ObjectName}}FromID").Arg("id", fields[i].Id)
+			{{- end }}
+			out = append(out, val)
+		}
+
+		return out
+	}
+	{{- end }}
+
+	{{- if and $field.TypeRef.IsList (IsListOfObject $field.TypeRef) }}
 	var response []{{ $field.Name | ToLowerCase }}
-    {{- else }}
+	{{- else }}
 	var response {{ $field.TypeRef | FormatOutputType }}
-    {{- end  }}
+	{{- end }}
 
 	q = q.Bind(&response)
 	{{- $typeName := $field.TypeRef | FormatOutputType }}
 	{{- if ne $typeName "Client" }}
-	    {{- if and $field.TypeRef.IsList (IsListOfObject $field.TypeRef) }}
+	{{- if and $field.TypeRef.IsList (IsListOfObject $field.TypeRef) }}
 
 	err := q.Execute(ctx)
 	if err != nil {
-	    return nil, err
+		return nil, err
 	}
 
 	return convert(response), nil
-	    {{- else }}
+	{{- else }}
 	return response, q.Execute(ctx)
-	    {{- end }}
+	{{- end }}
 	{{- else }}
 	return response, q.Execute(ctx, r.gql)
 	{{- end }}
@@ -195,30 +246,30 @@ func (r *{{ $.Name | FormatName }}) XXX_GraphQLIDType() string {
 
 // XXX_GraphQLID is an internal function. It returns the underlying type ID
 func (r *{{ $.Name | FormatName }}) XXX_GraphQLID(ctx context.Context) (string, error) {
-  id, err := r.ID(ctx)
-  if err != nil {
-    return "", err
-  }
+	id, err := r.ID(ctx)
+	if err != nil {
+		return "", err
+	}
 	return string(id), nil
 }
 
 func (r *{{ $.Name | FormatName }}) MarshalJSON() ([]byte, error) {
-  id, err := r.ID(marshalCtx)
-  if err != nil {
-    return nil, err
-  }
-  return json.Marshal(id)
+	id, err := r.ID(marshalCtx)
+	if err != nil {
+		return nil, err
+	}
+	return json.Marshal(id)
 }
 
 {{- if IsModuleCode }}
 func (r *{{ $.Name | FormatName }}) UnmarshalJSON(bs []byte) error {
-  var id string
-  err := json.Unmarshal(bs, &id)
-  if err != nil {
-    return err
-  }
-  *r = *dag.Load{{ $.Name | FormatName }}FromID({{$.Name | FormatName}}ID(id))
-  return nil
+	var id string
+	err := json.Unmarshal(bs, &id)
+	if err != nil {
+		return err
+	}
+	*r = *dag.Load{{ $.Name | FormatName }}FromID({{$.Name | FormatName}}ID(id))
+	return nil
 }
 {{- end }}
 

diff --git a/cmd/codegen/generator/typescript/templates/src/header.ts.gtpl b/cmd/codegen/generator/typescript/templates/src/header.ts.gtpl
@@ -8,11 +8,12 @@ inherited by futures objects and common types.
  * Do not make direct changes to the file.
  */
  {{- if IsBundle }}
-import { Context } from "./core.js"
+import { Context, runWithSpan } from "./core.js"
 {{- else if (not IsClientOnly)}}
 import { Context } from "../common/context.js"
+import { runWithSpan } from "../telemetry/index.js"
 {{- else }}
-import { Context, connect as _connect, connection as _connection, ConnectOpts, CallbackFct } from "@dagger.io/dagger"
+import { Context, connect as _connect, connection as _connection, ConnectOpts, CallbackFct, runWithSpan } from "@dagger.io/dagger"
 {{- end }}
 
 {{ if IsClientOnly }}
@@ -24,7 +25,7 @@ async function serveModuleDependencies(client: Client): Promise<void> {
   {{- range $i, $dep := $dependencies -}}
     {{- if eq $dep.Kind "GIT_SOURCE" }}
   await client.moduleSource(
-      "{{ $dep.Source }}", 
+      "{{ $dep.Source }}",
       { refPin: "{{ $dep.Pin }}" },
     )
     .withName("{{ $dep.Name }}")

diff --git a/cmd/codegen/generator/typescript/templates/src/object.ts.gtpl b/cmd/codegen/generator/typescript/templates/src/object.ts.gtpl
@@ -67,6 +67,27 @@ export class {{ .Name | QueryToClient | FormatName }} extends BaseClient { {{- w
 				{{- end }}
 			{{- end }}
 
+{{- if eq .Name "Status" }}
+
+  public async run<T>(fn: (span: Status) => Promise<T>) {
+    const started = await this.start()
+    const spanIdHex = await started.internalId()
+
+    let spanError: Error | undefined = undefined
+    try {
+      return await runWithSpan(fn, this, started, spanIdHex)
+    } catch (e: unknown) {
+      if (e instanceof globalThis.Error) {
+        spanError = dag.error(e.message)
+      } else {
+        spanError = dag.error(`Unknown error: ${e}`)
+      }
+      throw e
+    } finally {
+      await started.end({ error: spanError })
+    }
+  }
+{{- end }}
 {{- if . | IsSelfChainable }}
 {{""}}
   /**