probo/pkg/server/api/console/v1/agent_run_resolvers.go

package console_v1

// This file will be automatically regenerated based on the schema, any resolver
// implementations
// will be copied through when generating and any unknown code will be moved to the end.
// Code generated by github.com/99designs/gqlgen version v0.17.90

import (
	"context"
	"errors"
	"fmt"

	"github.com/vikstrous/dataloadgen"
	"go.gearno.de/kit/log"
	"go.probo.inc/probo/pkg/agent"
	"go.probo.inc/probo/pkg/agentrun"
	"go.probo.inc/probo/pkg/coredata"
	"go.probo.inc/probo/pkg/probo"
	"go.probo.inc/probo/pkg/server/api/console/v1/dataloader"
	"go.probo.inc/probo/pkg/server/api/console/v1/schema"
	"go.probo.inc/probo/pkg/server/api/console/v1/types"
	"go.probo.inc/probo/pkg/server/gqlutils"
)

// Organization is the resolver for the organization field.
func (r *agentRunResolver) Organization(ctx context.Context, obj *types.AgentRun) (*types.Organization, error) {
	if _, err := r.authorize(ctx, obj.ID, probo.ActionOrganizationGet); err != nil {
		return nil, err
	}

	loaders := dataloader.FromContext(ctx)

	organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
	if err != nil {
		if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
			return nil, gqlutils.NotFound(ctx, err)
		}

		r.logger.ErrorCtx(ctx, "cannot load organization", log.Error(err))

		return nil, gqlutils.Internal(ctx)
	}

	return types.NewOrganization(organization), nil
}

// Permission is the resolver for the permission field.
func (r *agentRunResolver) Permission(ctx context.Context, obj *types.AgentRun, action string) (bool, error) {
	return r.Resolver.Permission(ctx, obj, action)
}

// TotalCount is the resolver for the totalCount field.
func (r *agentRunConnectionResolver) TotalCount(ctx context.Context, obj *types.AgentRunConnection) (int, error) {
	scope, err := r.authorize(ctx, obj.ParentID, agentrun.ActionAgentRunList)
	if err != nil {
		return 0, err
	}

	switch obj.Resolver.(type) {
	case *organizationResolver:
		count, err := r.agentRun.CountForOrganizationID(ctx, scope, obj.ParentID)
		if err != nil {
			r.logger.ErrorCtx(ctx, "cannot count agent runs", log.Error(err))
			return 0, gqlutils.Internal(ctx)
		}

		return count, nil
	}

	r.logger.ErrorCtx(ctx, "unsupported resolver for agent run connection", log.String("resolver", fmt.Sprintf("%T", obj.Resolver)))

	return 0, gqlutils.Internal(ctx)
}

// SubmitAgentRunApproval is the resolver for the submitAgentRunApproval field.
func (r *mutationResolver) SubmitAgentRunApproval(ctx context.Context, input types.SubmitAgentRunApprovalInput) (*types.SubmitAgentRunApprovalPayload, error) {
	scope, err := r.authorize(ctx, input.AgentRunID, agentrun.ActionAgentRunApprove)
	if err != nil {
		return nil, err
	}

	decisions := make(map[string]agent.ApprovalResult, len(input.Decisions))
	for _, decision := range input.Decisions {
		message := ""
		if decision.Reason != nil {
			message = *decision.Reason
		}

		decisions[decision.ToolCallID] = agent.ApprovalResult{
			Approved: decision.Approved,
			Message:  message,
		}
	}

	run, err := r.agentRun.SubmitApproval(ctx, scope, input.AgentRunID, decisions)
	if err != nil {
		switch {
		case errors.Is(err, agentrun.ErrAgentRunNotFound):
			return nil, gqlutils.NotFound(ctx, err)
		case errors.Is(err, agentrun.ErrNotAwaitingApproval):
			return nil, gqlutils.Conflictf(ctx, "agent run is not awaiting approval")
		case errors.Is(err, agentrun.ErrApprovalDecisionsMismatch):
			return nil, gqlutils.Invalidf(ctx, "approval decisions must match the run's pending approvals")
		default:
			r.logger.ErrorCtx(ctx, "cannot submit agent run approval", log.Error(err))
			return nil, gqlutils.Internal(ctx)
		}
	}

	return &types.SubmitAgentRunApprovalPayload{
		AgentRun: types.NewAgentRun(run),
	}, nil
}

// AgentRun returns schema.AgentRunResolver implementation.
func (r *Resolver) AgentRun() schema.AgentRunResolver { return &agentRunResolver{r} }

// AgentRunConnection returns schema.AgentRunConnectionResolver implementation.
func (r *Resolver) AgentRunConnection() schema.AgentRunConnectionResolver {
	return &agentRunConnectionResolver{r}
}

type agentRunResolver struct{ *Resolver }
type agentRunConnectionResolver struct{ *Resolver }