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

package runtime

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"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
	}
	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)
		}
		body, manifest, err := projectcontext.Assemble(bundle, budget)
		if err != nil {
			return fail(err)
		}
		request := reminders.Attach(messages, []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"
		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
}