story/internal/workflowbridge/service.go

// Package workflowbridge executes private model steps for the Dify canvas.
// Elder speech and generated media live only in expiring memory, never in
// Dify's durable execution logs. The platform remains the authorized writer.
package workflowbridge

import (
	"context"
	"crypto/hmac"
	"crypto/rand"
	"data-terminal-cloud/internal/extraction"
	"data-terminal-cloud/internal/memoir"
	"encoding/hex"
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Editor interface {
	GenerateJSON(context.Context, string, string) (json.RawMessage, error)
	GenerateJSONSchema(context.Context, string, string, json.RawMessage) (json.RawMessage, error)
}
type ImageGenerator interface {
	Generate(context.Context, string) ([]byte, error)
}
type Service struct {
	Secret                string
	Extractor, BookEditor Editor
	Illustrator           ImageGenerator
	EditorForModel        func(step, model string) (Editor, error)
	ImageForModel         func(model string) (ImageGenerator, error)
	mu                    sync.Mutex
	contexts              map[string]*job
	imageBytes            int
}
type job struct {
	ID, Token, Kind string
	Expires         time.Time
	ctx             context.Context
	cancel          context.CancelFunc
	Input           extraction.GenerationInput
	Sources         []memoir.Source
	mu              sync.Mutex
	Result          extraction.GenerationResult
	Document        memoir.Document
	Pictures        map[int][]byte
	imageBytes      int
	ErrorCode       string
	Diagnostic      string
	Overdue         bool
	completed       map[string]map[string]any
	flights         map[string]chan struct{}
	models          map[string]string
}
type stepInput struct {
	ID          string `json:"context_id"`
	Token       string `json:"step_token"`
	Kind        string `json:"kind"`
	Prompt      string `json:"prompt"`
	Index       int    `json:"chapter_index"`
	Model       string `json:"model"`
	CustomModel string `json:"custom_model"`
}

func token() string { b := make([]byte, 24); _, _ = rand.Read(b); return hex.EncodeToString(b) }
func New(secret string, extractor, editor Editor, image ImageGenerator) *Service {
	return &Service{Secret: secret, Extractor: extractor, BookEditor: editor, Illustrator: image, contexts: map[string]*job{}}
}
func (s *Service) authorized(r *http.Request) bool {
	return s.Secret != "" && hmac.Equal([]byte(r.Header.Get("Authorization")), []byte("Bearer "+s.Secret))
}
func reply(w http.ResponseWriter, status int, out any) {
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Cache-Control", "no-store")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(out)
}
func (s *Service) create(kind string, raw json.RawMessage) (*job, error) {
	if kind != "extract" && kind != "memoir" {
		return nil, errors.New("invalid workflow kind")
	}
	c := &job{ID: token(), Token: token(), Kind: kind, Expires: time.Now().Add(20 * time.Minute), Pictures: map[int][]byte{}, completed: map[string]map[string]any{}, flights: map[string]chan struct{}{}}
	if kind == "extract" {
		if json.Unmarshal(raw, &c.Input) != nil || c.Input.Content == "" {
			return nil, errors.New("invalid extraction input")
		}
	} else {
		var in struct {
			Stories []memoir.Source `json:"stories"`
		}
		if json.Unmarshal(raw, &in) != nil || len(in.Stories) == 0 {
			return nil, errors.New("invalid book input")
		}
		c.Sources = in.Stories
	}
	c.ctx, c.cancel = context.WithTimeout(context.Background(), 20*time.Minute)
	s.mu.Lock()
	defer s.mu.Unlock()
	for id, old := range s.contexts {
		if time.Now().After(old.Expires) {
			old.cancel()
			s.imageBytes -= old.imageBytes
			delete(s.contexts, id)
		}
	}
	if len(s.contexts) >= 8 {
		c.cancel()
		return nil, errors.New("workflow capacity reached")
	}
	s.contexts[c.ID] = c
	time.AfterFunc(time.Until(c.Expires), func() { s.forget(c.ID) })
	return c, nil
}
func (s *Service) find(id string) *job {
	s.mu.Lock()
	defer s.mu.Unlock()
	c := s.contexts[id]
	if c != nil && time.Now().After(c.Expires) {
		c.cancel()
		delete(s.contexts, id)
		s.imageBytes -= c.imageBytes
		return nil
	}
	return c
}
func (s *Service) forget(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if c := s.contexts[id]; c != nil {
		c.cancel()
		delete(s.contexts, id)
		s.imageBytes -= c.imageBytes
	}
}

func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/health" {
		reply(w, 200, map[string]bool{"ready": true})
		return
	}
	if r.Method == "POST" && r.URL.Path == "/contexts" {
		if !s.authorized(r) {
			reply(w, 403, map[string]string{"error": "FORBIDDEN"})
			return
		}
		var in struct {
			Kind  string          `json:"kind"`
			Input json.RawMessage `json:"input"`
		}
		if json.NewDecoder(http.MaxBytesReader(w, r.Body, 2<<20)).Decode(&in) != nil {
			reply(w, 400, map[string]string{"error": "INVALID_INPUT"})
			return
		}
		c, err := s.create(in.Kind, in.Input)
		if err != nil {
			reply(w, 400, map[string]string{"error": "INVALID_INPUT"})
			return
		}
		reply(w, 200, map[string]string{"context_id": c.ID, "step_token": c.Token})
		return
	}
	if strings.HasPrefix(r.URL.Path, "/contexts/") {
		if !s.authorized(r) {
			reply(w, 403, map[string]string{"error": "FORBIDDEN"})
			return
		}
		parts := strings.Split(strings.Trim(r.URL.Path, "/"), "/")
		c := s.find(parts[1])
		if c == nil {
			reply(w, 404, map[string]string{"error": "CONTEXT_EXPIRED"})
			return
		}
		if r.Method == "DELETE" && len(parts) == 2 {
			s.forget(c.ID)
			reply(w, 200, map[string]bool{"deleted": true})
			return
		}
		c.mu.Lock()
		defer c.mu.Unlock()
		if r.Method == "GET" && len(parts) == 4 && parts[2] == "images" {
			i, err := strconv.Atoi(parts[3])
			b := c.Pictures[i]
			if err != nil || len(b) == 0 {
				w.WriteHeader(404)
				return
			}
			w.Header().Set("Content-Type", "image/jpeg")
			w.Header().Set("Cache-Control", "no-store")
			_, _ = w.Write(b)
			return
		}
		if r.Method == "GET" && len(parts) == 2 {
			doc := c.Document
			// Empty event arrays are meaningful: a today-only conversation has no
			// historical stories, and must survive the private JSON round trip.
			raw, _ := json.Marshal(c.Result)
			var result map[string]any
			_ = json.Unmarshal(raw, &result)
			result["stories"] = append([]extraction.EventStory{}, c.Result.Stories...)
			result["daily_events"] = append([]extraction.DailyEvent{}, c.Result.DailyEvents...)
			reply(w, 200, map[string]any{"result": result, "document": doc, "error_code": c.ErrorCode})
			return
		}
		if r.Method == "GET" && len(parts) == 3 && parts[2] == "diagnostic" {
			reply(w, 200, map[string]string{"model_json": c.Diagnostic})
			return
		}
	}
	if r.Method == "POST" && strings.HasPrefix(r.URL.Path, "/steps/") {
		var in stepInput
		if json.NewDecoder(http.MaxBytesReader(w, r.Body, 16<<10)).Decode(&in) != nil {
			reply(w, 400, map[string]string{"error": "INVALID_STEP"})
			return
		}
		step := strings.TrimPrefix(r.URL.Path, "/steps/")
		if step == "prepare" && in.ID == "" {
			var raw json.RawMessage
			if in.Kind == "extract" {
				raw, _ = json.Marshal(extraction.GenerationInput{TenantID: "workflow-demo", ChatDate: time.Now().Format("2006-01-02"), Content: "[demo-turn] 长者：今天吃了早饭。1970年9月我进厂学做手套，师傅先教我剪线头。", PreviousStories: []extraction.PreviousStory{}})
			} else {
				raw, _ = json.Marshal(map[string]any{"stories": []memoir.Source{{ID: "demo-story", Version: 1, Title: "进厂第一天", Content: "1970年9月我进厂学做手套，师傅先教我剪线头。", EventTime: "1970年9月"}}})
			}
			c, err := s.create(in.Kind, raw)
			if err != nil {
				reply(w, 429, map[string]string{"error": "CAPACITY_REACHED"})
				return
			}
			in.ID = c.ID
			in.Token = c.Token
		}
		c := s.find(in.ID)
		if c == nil || !hmac.Equal([]byte(in.Token), []byte(c.Token)) {
			reply(w, 403, map[string]string{"error": "STEP_CONTEXT_UNAVAILABLE"})
			return
		}
		if step == "prepare" {
			reply(w, 200, map[string]any{"success": true, "context_id": c.ID, "step_token": c.Token, "kind": c.Kind})
			return
		}
		out := s.step(c, step, in)
		reply(w, 200, out)
		return
	}
	reply(w, 404, map[string]string{"error": "NOT_FOUND"})
}

func (s *Service) step(c *job, step string, in stepInput) map[string]any {
	key := step
	if step == "illustrate" {
		key += ":" + strconv.Itoa(in.Index)
	}
	for {
		c.mu.Lock()
		if out, ok := c.completed[key]; ok {
			c.mu.Unlock()
			return out
		}
		if flight := c.flights[key]; flight != nil {
			c.mu.Unlock()
			select {
			case <-flight:
				continue
			case <-c.ctx.Done():
				return map[string]any{"success": false, "error_code": "WORKFLOW_CANCELLED"}
			}
		}
		flight := make(chan struct{})
		c.flights[key] = flight
		overdue := c.Overdue
		c.mu.Unlock()
		var out map[string]any
		var err error
		if overdue && (step == "classify" || step == "daily" || step == "write-book" || step == "illustrate") {
			err = extraction.ErrModelAccountOverdue
		} else {
			out, err = s.execute(c, step, in)
		}
		c.mu.Lock()
		if err != nil {
			code := "MODEL_STEP_FAILED"
			if errors.Is(err, errInvalidModel) {
				code = "INVALID_MODEL_SELECTION"
			}
			if errors.Is(err, extraction.ErrModelAccountOverdue) {
				code = "MODEL_ACCOUNT_OVERDUE"
				c.Overdue = true
			}
			c.ErrorCode = code
			out = map[string]any{"success": false, "error_code": code, "step": step}
			if step == "illustrate" {
				c.Document.IllustrationErrors++
				out["chapter_index"] = in.Index
			}
		}
		c.completed[key] = out
		delete(c.flights, key)
		close(flight)
		c.mu.Unlock()
		return out
	}
}
func (s *Service) execute(c *job, step string, in stepInput) (map[string]any, error) {
	model, err := selectedModel(step, in.Model, in.CustomModel)
	if err != nil {
		return nil, err
	}
	extractor, editor, image := s.Extractor, s.BookEditor, s.Illustrator
	if model != "" {
		if step == "illustrate" && s.ImageForModel != nil {
			image, err = s.ImageForModel(model)
		} else if s.EditorForModel != nil {
			var selected Editor
			selected, err = s.EditorForModel(step, model)
			if step == "write-book" {
				editor = selected
			} else {
				extractor = selected
			}
		}
		if err != nil {
			return nil, err
		}
		c.mu.Lock()
		if c.models == nil {
			c.models = map[string]string{}
		}
		c.models[step] = model
		c.mu.Unlock()
	}
	classified, daily := extraction.WorkflowPrompts()
	switch step {
	case "classify":
		if c.Kind != "extract" {
			break
		}
		prompt := in.Prompt
		if prompt == "" {
			prompt = classified
		}
		raw, err := extractor.GenerateJSON(c.ctx, extraction.WorkflowContent(c.Input), prompt)
		if err != nil {
			return nil, err
		}
		var result extraction.GenerationResult
		if json.Unmarshal(raw, &result) != nil || result.Stories == nil {
			return nil, errors.New("invalid extraction result")
		}
		c.mu.Lock()
		c.Result = result
		c.mu.Unlock()
		return map[string]any{"success": true, "model": model, "story_count": len(result.Stories), "daily_candidate_count": len(result.DailyEvents)}, nil
	case "daily":
		if c.Kind != "extract" {
			break
		}
		prompt := in.Prompt
		if prompt == "" {
			prompt = daily
		}
		input, _ := json.Marshal(map[string]any{"chat_date": c.Input.ChatDate, "transcript": c.Input.Content})
		raw, err := extractor.GenerateJSON(c.ctx, string(input), prompt)
		if err != nil {
			return nil, err
		}
		var result struct {
			Events []extraction.DailyEvent `json:"daily_events"`
		}
		if json.Unmarshal(raw, &result) != nil || result.Events == nil {
			return nil, errors.New("invalid daily result")
		}
		c.mu.Lock()
		c.Result.DailyEvents = result.Events
		c.mu.Unlock()
		return map[string]any{"success": true, "model": model, "today_actual_count": len(result.Events)}, nil
	case "dedup":
		if c.Kind != "extract" {
			break
		}
		c.mu.Lock()
		defer c.mu.Unlock()
		r, err := extraction.ValidateWorkflowResult(c.Input, c.Result, "dify:"+c.models["classify"])
		if err != nil {
			return nil, err
		}
		c.Result = r
		created, updated, repeated := 0, 0, 0
		for _, story := range r.Stories {
			if story.MatchID == "" {
				created++
			} else if len(story.NewDetails) > 0 {
				updated++
			} else {
				repeated++
			}
		}
		return map[string]any{"success": true, "story_count": len(r.Stories), "daily_count": len(r.DailyEvents), "new_stories": created, "updated_stories": updated, "repeated_stories": repeated}, nil
	case "write-book":
		if c.Kind != "memoir" {
			break
		}
		prompt := in.Prompt
		if prompt == "" {
			prompt = memoir.MemoirPrompt
		}
		input, _ := json.Marshal(map[string]any{"stories": c.Sources})
		var doc memoir.Document
		var err error
		// Bound all attempts inside the HTTP node's 660-second read budget.
		// Whole-book prose can take longer than a short chat classification.
		stepCtx, cancel := context.WithTimeout(c.ctx, 600*time.Second)
		defer cancel()
		for attempt := 0; attempt < 3; attempt++ {
			var raw json.RawMessage
			raw, err = editor.GenerateJSONSchema(stepCtx, string(input), prompt, json.RawMessage(memoir.MemoirSchema))
			if err != nil && len(raw) > 0 {
				c.mu.Lock()
				c.Diagnostic = string(raw)
				c.mu.Unlock()
			}
			if errors.Is(err, extraction.ErrModelAccountOverdue) {
				return nil, err
			}
			if err == nil {
				err = json.Unmarshal(raw, &doc)
				if err == nil {
					err = memoir.ValidateDocument(doc, c.Sources)
				}
				if err == nil {
					break
				}
			}
			// Errors contain fixed categories or byte counts, never model prose.
			log.Printf("workflow book attempt=%d error=%v", attempt+1, err)
			if stepCtx.Err() != nil {
				break
			}
		}
		if err != nil {
			return nil, err
		}
		ordered := make([]memoir.Chapter, 0, len(c.Sources))
		for _, source := range c.Sources {
			for _, chapter := range doc.Chapters {
				if chapter.StoryID == source.ID {
					ordered = append(ordered, chapter)
					break
				}
			}
		}
		doc.Chapters = ordered
		c.mu.Lock()
		c.Document = doc
		c.mu.Unlock()
		indexes := make([]int, len(doc.Chapters))
		for i := range indexes {
			indexes[i] = i
		}
		return map[string]any{"success": true, "model": model, "chapter_count": len(indexes), "chapter_indexes": indexes}, nil
	case "validate-book":
		if c.Kind != "memoir" {
			break
		}
		c.mu.Lock()
		defer c.mu.Unlock()
		if err := memoir.ValidateDocument(c.Document, c.Sources); err != nil {
			return nil, err
		}
		return map[string]any{"success": true, "chapter_count": len(c.Document.Chapters), "all_stories_covered": true}, nil
	case "illustrate":
		if c.Kind != "memoir" {
			break
		}
		c.mu.Lock()
		if in.Index < 0 || in.Index >= len(c.Document.Chapters) {
			c.mu.Unlock()
			break
		}
		prompt := c.Document.Chapters[in.Index].IllustrationPrompt
		c.mu.Unlock()
		b, err := image.Generate(c.ctx, prompt)
		if err != nil {
			return nil, err
		}
		s.mu.Lock()
		if s.contexts[c.ID] != c || s.imageBytes+len(b) > 192<<20 {
			s.mu.Unlock()
			return nil, errors.New("image memory limit")
		}
		s.imageBytes += len(b)
		c.imageBytes += len(b)
		s.mu.Unlock()
		c.mu.Lock()
		c.Pictures[in.Index] = b
		c.mu.Unlock()
		return map[string]any{"success": true, "model": model, "chapter_index": in.Index, "illustration_ready": true}, nil
	case "finish":
		c.mu.Lock()
		defer c.mu.Unlock()
		return map[string]any{"success": true, "story_count": len(c.Result.Stories), "chapter_count": len(c.Document.Chapters), "illustrations_ready": len(c.Pictures), "illustration_errors": c.Document.IllustrationErrors}, nil
	}
	return nil, errors.New("invalid step for workflow context")
}