Assistant/cmd/agent-worker/task_evaluation_executor.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"git.cin.su/m/Assistant/contracts/artifact"
	"git.cin.su/m/Assistant/contracts/capability"
	"git.cin.su/m/Assistant/contracts/work"
	capabilitydomain "git.cin.su/m/Assistant/modules/capability/domain"
)

type taskEvaluationExecutor struct {
	cognitiveURL string
	callbackURL  string
	verifierURL  string
	client       *http.Client
	fallback     localExecutor
}

func newTaskEvaluationExecutor(cognitiveURL, callbackURL, verifierURL string, client *http.Client) taskEvaluationExecutor {
	if client == nil {
		client = http.DefaultClient
	}
	return taskEvaluationExecutor{cognitiveURL: strings.TrimRight(cognitiveURL, "/"), callbackURL: strings.TrimRight(callbackURL, "/"), verifierURL: strings.TrimRight(verifierURL, "/"), client: client, fallback: newLocalExecutor()}
}

func (e taskEvaluationExecutor) Execute(ctx context.Context, request work.AgentExecutionRequest) (work.AgentExecutionResult, error) {
	intent := request.Intent
	if intent.TaskEvaluationExecution == nil || intent.TaskEvaluationProfile == nil || len(intent.TaskEvaluationCases) == 0 {
		return e.fallback.Execute(ctx, request)
	}
	execution := intent.TaskEvaluationExecution
	if execution.GraphID == "" || execution.NodeID == "" || execution.BindingRevision < 1 || execution.ActorID == "" || execution.PrincipalID == "" || execution.ScopeID == "" {
		return e.fallback.Execute(ctx, request)
	}
	if e.cognitiveURL == "" || e.callbackURL == "" || e.verifierURL == "" {
		return e.fallback.Execute(ctx, request)
	}
	payload := capability.TaskEvaluationExecutionRequest{Profile: *intent.TaskEvaluationProfile, CapabilityID: intent.TaskEvaluationProfile.CapabilityID, BindingID: intent.TaskEvaluationBindingID, BindingRevision: execution.BindingRevision, ActorID: execution.ActorID, PrincipalID: execution.PrincipalID, ScopeID: execution.ScopeID, Cases: intent.TaskEvaluationCases, MaxOutputBytes: 4096}
	response, err := postJSON[capability.TaskEvaluationExecutionResponse](ctx, e.client, e.cognitiveURL+"/v1/task-evaluations/execute", payload)
	if err != nil {
		return work.AgentExecutionResult{}, err
	}
	if len(response.Responses) == 0 || response.ObservedAtUnixMillis < 1 {
		return work.AgentExecutionResult{}, fmt.Errorf("task evaluation returned no observations")
	}
	if len(response.Observations) == len(intent.TaskEvaluationCases) && taskEvaluationHasCriticalFailure(response.Observations) {
		evaluation, err := capabilitydomain.BuildTaskEvaluation(*intent.TaskEvaluationProfile, capability.TaskEvaluationID("evaluation:"+intent.DispatchID), response.BindingID, response.BindingRevision, intent.TaskEvaluationProfile.CapabilityID, response.Observations, artifact.ArtifactID(intent.TaskEvaluationCases[0].EvidenceArtifactRef), response.ObservedAtUnixMillis)
		if err != nil {
			return work.AgentExecutionResult{}, err
		}
		callback := work.TaskEvaluationEvidenceCallback{GraphID: execution.GraphID, Revision: execution.GraphRevision, NodeID: execution.NodeID, Evaluation: evaluation, Succeeded: true}
		if _, err := postJSON[json.RawMessage](ctx, e.client, e.callbackURL+"/v1/work-graphs/task-evaluations/callback", callback); err != nil {
			return work.AgentExecutionResult{}, err
		}
		return work.AgentExecutionResult{State: work.AgentExecutionCompleted, Summary: "task evaluation recorded as UNKNOWN after bounded provider recovery", OutcomeRef: string(evaluation.ID), ArtifactRefs: []string{string(evaluation.ArtifactID)}, ObservedAtUnixMillis: time.Now().UnixMilli()}, nil
	}
	if len(intent.TaskEvaluationCases) == 0 {
		return work.AgentExecutionResult{}, fmt.Errorf("task evaluation cases are missing")
	}
	caseByID := make(map[string]capability.TaskEvaluationCase, len(intent.TaskEvaluationCases))
	for _, item := range intent.TaskEvaluationCases {
		caseByID[item.CaseID] = item
	}
	observations := make([]capability.TaskCaseObservation, 0, len(response.Responses))
	for _, item := range response.Responses {
		caseValue, ok := caseByID[item.CaseID]
		if !ok || caseValue.Verification.Mode == "" {
			return e.fallback.Execute(ctx, request)
		}
		verified, verifyErr := postJSON[capability.TaskCaseVerificationResponse](ctx, e.client, e.verifierURL+"/v1/task-evaluations/verify", capability.TaskCaseVerificationRequest{CaseID: item.CaseID, ResponseText: item.ResponseText, LatencyMicros: item.LatencyMicros, EvidenceArtifactRef: item.EvidenceArtifactRef, Verification: caseValue.Verification})
		if verifyErr != nil {
			return e.fallback.Execute(ctx, request)
		}
		observations = append(observations, capability.TaskCaseObservation{CaseID: verified.CaseID, Result: verified.Result, QualityScoreMillis: verified.QualityScoreMillis, LatencyMicros: verified.LatencyMicros, CriticalFailure: verified.CriticalFailure, EvidenceArtifactRef: verified.EvidenceArtifactRef, VerificationEvidenceRef: verified.VerificationEvidenceRef, VerifierIdentity: verified.VerifierIdentity})
	}
	evaluation, err := capabilitydomain.BuildTaskEvaluation(*intent.TaskEvaluationProfile, capability.TaskEvaluationID("evaluation:"+intent.DispatchID), response.BindingID, response.BindingRevision, intent.TaskEvaluationProfile.CapabilityID, observations, artifact.ArtifactID(intent.TaskEvaluationCases[0].EvidenceArtifactRef), response.ObservedAtUnixMillis)
	if err != nil {
		return work.AgentExecutionResult{}, err
	}
	callback := work.TaskEvaluationEvidenceCallback{GraphID: execution.GraphID, Revision: execution.GraphRevision, NodeID: execution.NodeID, Evaluation: evaluation, Succeeded: true}
	if _, err := postJSON[json.RawMessage](ctx, e.client, e.callbackURL+"/v1/work-graphs/task-evaluations/callback", callback); err != nil {
		return work.AgentExecutionResult{}, err
	}
	return work.AgentExecutionResult{State: work.AgentExecutionCompleted, Summary: "task evaluation executed, aggregated and accepted by Work Graph callback", OutcomeRef: string(evaluation.ID), ArtifactRefs: []string{string(evaluation.ArtifactID)}, ObservedAtUnixMillis: time.Now().UnixMilli()}, nil
}

func taskEvaluationHasCriticalFailure(observations []capability.TaskCaseObservation) bool {
	for _, observation := range observations {
		if observation.CriticalFailure {
			return true
		}
	}
	return false
}

func postJSON[T any](ctx context.Context, client *http.Client, url string, payload any) (T, error) {
	var result T
	body, err := json.Marshal(payload)
	if err != nil {
		return result, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return result, err
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := client.Do(req)
	if err != nil {
		return result, err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return result, fmt.Errorf("POST %s: status %s", url, resp.Status)
	}
	if resp.StatusCode == http.StatusNoContent {
		return result, nil
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return result, err
	}
	return result, nil
}