Agent-System-Backend/internal/runtime/project_stream.go

package runtime

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	agentllm "agent-native/runtime/llm"
	"agent-native/runtime/reminders"
	"agent-system-backend/internal/ctxasm"
	"agent-system-backend/internal/projectcontext"
	"agent-system-backend/pkg/extension"
	core "agent-system-core"
)

// GuardProjectStream injects project reference material only into a model
// request copy. This is used by execution paths that do not use LLMRunner's
// direct loop; no project or no registered provider keeps the old request.
func GuardProjectStream(rc any, goal string, stream agentllm.StreamFunc) agentllm.StreamFunc {
	source, ok := rc.(interface {
		ProjectContext(context.Context, string, int) (extension.ProjectContextBundle, bool, error)
	})
	if stream == nil {
		return stream
	}
	if !ok {
		if scope, hasRun := rc.(interface{ Run() core.Run }); hasRun && scope.Run().ProjectID != "" {
			return func(context.Context, string, []agentllm.Message, []agentllm.ToolDecl, agentllm.StreamOptions) <-chan agentllm.Event {
				out := make(chan agentllm.Event, 1)
				out <- agentllm.Event{Kind: agentllm.KindError, Err: fmt.Errorf("trusted project context port missing")}
				close(out)
				return out
			}
		}
		return stream
	}
	var scopeMu sync.Mutex
	var previousScope string
	var seenScope, taintedHistory bool
	var cleanPrefixIndex int
	var cleanTask agentllm.Message
	var droppedPriorHistory bool
	return func(ctx context.Context, model string, messages []agentllm.Message, tools []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		fail := func(err error) <-chan agentllm.Event {
			out := make(chan agentllm.Event, 1)
			out <- agentllm.Event{Kind: agentllm.KindError, Err: err}
			close(out)
			return out
		}
		const budget = 12000
		requestGoal := goal
		if requestGoal == "" {
			for i := len(messages) - 1; i >= 0; i-- {
				if messages[i].Role == agentllm.RoleUser {
					requestGoal = messages[i].Content
					break
				}
			}
		}
		bundle, active, err := source.ProjectContext(ctx, requestGoal, budget)
		if err != nil {
			return fail(err)
		}
		if !active {
			return stream(ctx, model, messages, tools, opts)
		}
		if checker, ok := rc.(interface {
			ValidateProjectCrossGuard(context.Context, extension.ProjectContextBundle) error
		}); ok {
			if err := checker.ValidateProjectCrossGuard(ctx, bundle); err != nil {
				return fail(err)
			}
		}
		body, manifest, err := projectcontext.Assemble(bundle, budget)
		if err != nil {
			return fail(err)
		}
		scopeMu.Lock()
		scope := projectScopeKey(manifest)
		requestBase := messages
		if !seenScope {
			users, lastUser, lastMachine := 0, -1, -1
			for i, msg := range messages {
				if msg.Role == agentllm.RoleAssistant || msg.Role == agentllm.RoleTool {
					lastMachine = i
				}
				if msg.Role == agentllm.RoleUser {
					users++
					lastUser = i
				}
			}
			if lastUser < 0 || lastUser <= lastMachine {
				scopeMu.Unlock()
				return fail(fmt.Errorf("project model request lacks a fresh user task after unproven machine history"))
			}
			if lastMachine >= 0 || users != 1 {
				// This backend may have loaded a prior session. Keep only the new
				// task, then attach the freshly resolved bundle below.
				droppedPriorHistory, cleanPrefixIndex, cleanTask = true, lastUser, messages[lastUser]
			}
		}
		if seenScope && (scope != previousScope || manifest.AuthorizationEpoch == "" || manifest.PublicationEpoch == "") {
			taintedHistory = true
		}
		previousScope, seenScope = scope, true
		unsafe := taintedHistory
		if droppedPriorHistory && !unsafe {
			requestBase = []agentllm.Message{{Role: agentllm.RoleSystem, Content: "Work on the current authorized project task. Prior machine history was discarded."}}
			if cleanPrefixIndex < len(messages) && messages[cleanPrefixIndex].Role == agentllm.RoleUser && messages[cleanPrefixIndex].Content == cleanTask.Content {
				requestBase = append(requestBase, messages[cleanPrefixIndex:]...)
			} else {
				// A compacted or rewritten session lost the trusted boundary.
				// Continue from the saved user task without prior machine text.
				requestBase = append(requestBase, cleanTask)
			}
		}
		scopeMu.Unlock()
		if unsafe {
			return fail(fmt.Errorf("project authorization changed; generic model history lacks source provenance"))
		}
		request := reminders.Attach(requestBase, []string{body})
		if err := checkProjectInputBudget(rc, request, tools); err != nil {
			return fail(err)
		}
		wire, err := json.Marshal(request)
		if err != nil {
			return fail(err)
		}
		manifest.Model = model
		manifest.InputSHA256 = fmt.Sprintf("%x", sha256.Sum256(wire))
		if ledger, ok := rc.(interface {
			ContextEvidenceInRequest([]agentllm.Message) []extension.ProjectContextManifestEntry
		}); ok {
			manifest.ReadEvidence = ledger.ContextEvidenceInRequest(request)
		}
		attemptID, err := projectcontext.NewAttemptID()
		if err != nil {
			return fail(err)
		}
		manifest.AttemptID, manifest.Delivery = attemptID, "attempted"
		manifest.WorkflowStepAttemptID = projectcontext.WorkflowStepAttempt(ctx)
		var save func(context.Context, extension.ProjectContextManifest) error
		if sink, ok := rc.(interface {
			SaveContextManifest(context.Context, extension.ProjectContextManifest) error
		}); ok {
			save = sink.SaveContextManifest
			if err := save(ctx, manifest); err != nil {
				return fail(err)
			}
		}
		in := stream(ctx, model, request, tools, opts)
		out := make(chan agentllm.Event, 1)
		go func() {
			defer close(out)
			terminal := false
			finish := func(status string) error {
				if save == nil {
					return nil
				}
				manifest.Delivery = status
				receiptCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
				defer cancel()
				return save(receiptCtx, manifest)
			}
			defer func() {
				if !terminal {
					_ = finish("unknown")
				}
			}()
			for ev := range in {
				if ev.Kind == agentllm.KindMessageEnd || ev.Kind == agentllm.KindError {
					status := "unknown"
					if ev.Kind == agentllm.KindMessageEnd {
						status = "responded"
					}
					if err := finish(status); err != nil {
						select {
						case out <- agentllm.Event{Kind: agentllm.KindError, Err: err}:
						case <-ctx.Done():
						}
						terminal = true
						return
					}
					terminal = true
				}
				select {
				case out <- ev:
				case <-ctx.Done():
					return
				}
			}
		}()
		return out
	}
}

func checkProjectInputBudget(rc any, messages []agentllm.Message, tools []agentllm.ToolDecl) error {
	settings, ok := rc.(interface{ Settings() core.EffectiveSettings })
	if !ok {
		return fmt.Errorf("project model budget unavailable")
	}
	eff := settings.Settings()
	limit := int(eff.GetNumber("models.chat.contextWindow", 128000) -
		eff.GetNumber("models.chat.outputReserve", 16384) -
		eff.GetNumber("models.chat.contextBuffer", 4096))
	if limit <= 0 {
		return fmt.Errorf("project model input budget exhausted")
	}
	used := 0
	for _, msg := range messages {
		used += ctxasm.EstimateTokens(msg.Content) + 8
	}
	for _, tool := range tools {
		used += ctxasm.EstimateTokens(tool.Description+string(tool.Parameters)) + 8
	}
	if used > limit {
		return fmt.Errorf("project model input exceeds budget: estimated %d tokens, available %d", used, limit)
	}
	return nil
}

func projectScopeKey(m extension.ProjectContextManifest) string {
	selected, _ := json.Marshal(struct {
		Included []extension.ProjectContextManifestEntry
		Offered  []extension.ProjectEvidenceRef
		Missing  []string
	}{m.Included, m.OfferedReferences, m.Missing})
	return fmt.Sprintf("%s\x00%s\x00%s\x00%s\x00%s\x00%s", m.ProjectID, m.Version, m.Watermark, m.AuthorizationEpoch, m.PublicationEpoch, selected)
}