perf: batch workflow prompt and profile meta fetch (#2377)
Co-authored-by: Carlos Florencio <carlosmflorencio@gmail.com>
This commit is contained in:
@@ -393,14 +393,16 @@ func (a *orchestratorWorkflowStepGetterAdapter) GetPreviousStepByPosition(ctx co
|
||||
return a.svc.GetPreviousStepByPosition(ctx, workflowID, currentPosition)
|
||||
}
|
||||
|
||||
// GetWorkflowAgentProfileID implements orchestrator.WorkflowStepGetter.
|
||||
func (a *orchestratorWorkflowStepGetterAdapter) GetWorkflowAgentProfileID(ctx context.Context, workflowID string) (string, error) {
|
||||
return a.svc.GetWorkflowAgentProfileID(ctx, workflowID)
|
||||
}
|
||||
|
||||
// GetWorkflowPrompt implements orchestrator.WorkflowStepGetter.
|
||||
func (a *orchestratorWorkflowStepGetterAdapter) GetWorkflowPrompt(ctx context.Context, workflowID string) (string, error) {
|
||||
return a.svc.GetWorkflowPrompt(ctx, workflowID)
|
||||
// GetWorkflowMeta implements orchestrator.WorkflowStepGetter.
|
||||
func (a *orchestratorWorkflowStepGetterAdapter) GetWorkflowMeta(ctx context.Context, workflowID string) (orchestrator.WorkflowMeta, error) {
|
||||
meta, err := a.svc.GetWorkflowMeta(ctx, workflowID)
|
||||
if err != nil {
|
||||
return orchestrator.WorkflowMeta{}, err
|
||||
}
|
||||
return orchestrator.WorkflowMeta{
|
||||
AgentProfileID: meta.AgentProfileID,
|
||||
Prompt: meta.Prompt,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// reviewTaskCreatorAdapter adapts the task service to the orchestrator's ReviewTaskCreator interface.
|
||||
|
||||
@@ -40,8 +40,12 @@ import (
|
||||
// mockStepGetter implements WorkflowStepGetter for testing.
|
||||
type mockStepGetter struct {
|
||||
steps map[string]*wfmodels.WorkflowStep // stepID -> step
|
||||
workflowAgentProfileID string // returned by GetWorkflowAgentProfileID
|
||||
workflowAgentProfileID string // returned by GetWorkflowMeta
|
||||
workflowPrompts map[string]string // workflowID -> prompt
|
||||
workflowMetaCalls int // GetWorkflowMeta invocations
|
||||
workflowMetaErr error // optional error from GetWorkflowMeta
|
||||
workflowMetaDelay time.Duration // optional sleep before returning meta
|
||||
workflowMetaMu sync.Mutex // guards workflowMetaCalls for concurrent tests
|
||||
}
|
||||
|
||||
func newMockStepGetter() *mockStepGetter {
|
||||
@@ -82,18 +86,31 @@ func (m *mockStepGetter) GetPreviousStepByPosition(_ context.Context, workflowID
|
||||
return best, nil
|
||||
}
|
||||
|
||||
func (m *mockStepGetter) GetWorkflowAgentProfileID(_ context.Context, workflowID string) (string, error) {
|
||||
if m.workflowAgentProfileID != "" {
|
||||
return m.workflowAgentProfileID, nil
|
||||
func (m *mockStepGetter) GetWorkflowMeta(_ context.Context, workflowID string) (WorkflowMeta, error) {
|
||||
m.workflowMetaMu.Lock()
|
||||
m.workflowMetaCalls++
|
||||
delay := m.workflowMetaDelay
|
||||
m.workflowMetaMu.Unlock()
|
||||
if delay > 0 {
|
||||
time.Sleep(delay)
|
||||
}
|
||||
return "", nil
|
||||
if m.workflowMetaErr != nil {
|
||||
return WorkflowMeta{}, m.workflowMetaErr
|
||||
}
|
||||
prompt := ""
|
||||
if m.workflowPrompts != nil {
|
||||
prompt = m.workflowPrompts[workflowID]
|
||||
}
|
||||
return WorkflowMeta{
|
||||
AgentProfileID: m.workflowAgentProfileID,
|
||||
Prompt: prompt,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (m *mockStepGetter) GetWorkflowPrompt(_ context.Context, workflowID string) (string, error) {
|
||||
if m.workflowPrompts == nil {
|
||||
return "", nil
|
||||
}
|
||||
return m.workflowPrompts[workflowID], nil
|
||||
func (m *mockStepGetter) metaCalls() int {
|
||||
m.workflowMetaMu.Lock()
|
||||
defer m.workflowMetaMu.Unlock()
|
||||
return m.workflowMetaCalls
|
||||
}
|
||||
|
||||
// mockTaskRepo implements scheduler.TaskRepository for testing.
|
||||
|
||||
@@ -763,14 +763,14 @@ func (s *Service) resolveStepAgentProfile(ctx context.Context, step *wfmodels.Wo
|
||||
return step.AgentProfileID
|
||||
}
|
||||
if s.workflowStepGetter != nil && step.WorkflowID != "" {
|
||||
wfProfileID, err := s.workflowStepGetter.GetWorkflowAgentProfileID(ctx, step.WorkflowID)
|
||||
meta, err := s.getWorkflowMeta(ctx, step.WorkflowID)
|
||||
if err != nil {
|
||||
s.logger.Warn("failed to resolve workflow agent profile, falling back to task defaults",
|
||||
zap.String("workflow_id", step.WorkflowID),
|
||||
zap.String("step_id", step.ID),
|
||||
zap.Error(err))
|
||||
} else if wfProfileID != "" {
|
||||
return wfProfileID
|
||||
} else if meta.AgentProfileID != "" {
|
||||
return meta.AgentProfileID
|
||||
}
|
||||
}
|
||||
return ""
|
||||
@@ -1077,6 +1077,9 @@ func (s *Service) maybySwitchSessionForProfile(
|
||||
|
||||
// processOnEnter processes the on_enter events for a step after transitioning to it.
|
||||
func (s *Service) processOnEnter(ctx context.Context, taskID string, session *models.TaskSession, step *wfmodels.WorkflowStep, taskDescription string) {
|
||||
// One GetWorkflowMeta read shared by profile resolution and prompt build.
|
||||
ctx = withWorkflowMetaCache(ctx)
|
||||
|
||||
// Switch session if this step requires a different agent profile.
|
||||
var ok bool
|
||||
prevSessionID := session.ID
|
||||
|
||||
@@ -145,15 +145,21 @@ type taskQueuePromotionPublisher interface {
|
||||
PublishTaskQueuePromoted(ctx context.Context, task *models.Task)
|
||||
}
|
||||
|
||||
// WorkflowMeta is the subset of workflow fields needed at step entry
|
||||
// (agent profile default + optional workflow-level prompt).
|
||||
type WorkflowMeta struct {
|
||||
AgentProfileID string
|
||||
Prompt string
|
||||
}
|
||||
|
||||
// WorkflowStepGetter retrieves workflow step information for prompt building.
|
||||
type WorkflowStepGetter interface {
|
||||
GetStep(ctx context.Context, stepID string) (*wfmodels.WorkflowStep, error)
|
||||
GetNextStepByPosition(ctx context.Context, workflowID string, currentPosition int) (*wfmodels.WorkflowStep, error)
|
||||
GetPreviousStepByPosition(ctx context.Context, workflowID string, currentPosition int) (*wfmodels.WorkflowStep, error)
|
||||
GetWorkflowAgentProfileID(ctx context.Context, workflowID string) (string, error)
|
||||
// GetWorkflowPrompt returns the optional workflow-level agent instructions
|
||||
// prepended at step entry. Empty string means the workflow has no prompt.
|
||||
GetWorkflowPrompt(ctx context.Context, workflowID string) (string, error)
|
||||
// GetWorkflowMeta returns agent profile id and prompt for a workflow in one
|
||||
// read. Step-entry paths that need both fields should call this once.
|
||||
GetWorkflowMeta(ctx context.Context, workflowID string) (WorkflowMeta, error)
|
||||
}
|
||||
|
||||
// PromptReferenceExpander resolves "@name" saved-prompt references embedded in
|
||||
|
||||
@@ -337,6 +337,9 @@ func (s *Service) StartCreatedSession(
|
||||
attachments []v1.MessageAttachment,
|
||||
references []v1.EntityReference,
|
||||
) (*executor.TaskExecution, error) {
|
||||
// One GetWorkflowMeta read shared by profile resolution and prompt build.
|
||||
ctx = withWorkflowMetaCache(ctx)
|
||||
|
||||
s.logger.Debug("starting created session",
|
||||
zap.String("task_id", taskID),
|
||||
zap.String("session_id", sessionID),
|
||||
@@ -778,6 +781,9 @@ func (s *Service) StartTaskWithRoute(
|
||||
|
||||
//nolint:cyclop,funlen,gocognit // launch path threads many orthogonal concerns (workflow-step / agent-profile / office-task / config-mode / route / system-prompt wrapping); splitting it would require shared mutable state across helpers
|
||||
func (s *Service) startTask(ctx context.Context, taskID string, agentProfileID string, executorID string, executorProfileID string, priority string, prompt string, workflowStepID string, planMode, autoStart bool, attachments []v1.MessageAttachment, opts startTaskOptions) (*executor.TaskExecution, error) {
|
||||
// One GetWorkflowMeta read shared by profile resolution and prompt build.
|
||||
ctx = withWorkflowMetaCache(ctx)
|
||||
|
||||
env, route := opts.Env, opts.Route
|
||||
s.logger.Debug("manually starting task",
|
||||
zap.String("task_id", taskID),
|
||||
@@ -1494,7 +1500,7 @@ func (s *Service) workflowInstructionsBlock(ctx context.Context, step *wfmodels.
|
||||
if s.workflowStepGetter == nil || step == nil || step.WorkflowID == "" {
|
||||
return ""
|
||||
}
|
||||
prompt, err := s.workflowStepGetter.GetWorkflowPrompt(ctx, step.WorkflowID)
|
||||
meta, err := s.getWorkflowMeta(ctx, step.WorkflowID)
|
||||
if err != nil {
|
||||
if s.logger != nil {
|
||||
s.logger.Warn("failed to get workflow prompt for prompt building",
|
||||
@@ -1503,7 +1509,7 @@ func (s *Service) workflowInstructionsBlock(ctx context.Context, step *wfmodels.
|
||||
}
|
||||
return ""
|
||||
}
|
||||
prompt = strings.TrimSpace(prompt)
|
||||
prompt := strings.TrimSpace(meta.Prompt)
|
||||
if prompt == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
package orchestrator
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
"golang.org/x/sync/singleflight"
|
||||
)
|
||||
|
||||
// workflowMetaCacheKey is the context key for a request-scoped WorkflowMeta cache.
|
||||
// Step-entry paths seed the cache once so profile resolution and prompt build
|
||||
// share a single GetWorkflowMeta provider read.
|
||||
type workflowMetaCacheKey struct{}
|
||||
|
||||
type workflowMetaCacheEntry struct {
|
||||
meta WorkflowMeta
|
||||
err error
|
||||
}
|
||||
|
||||
type workflowMetaCache struct {
|
||||
mu sync.Mutex
|
||||
byID map[string]workflowMetaCacheEntry
|
||||
// group coalesces concurrent same-ID loads within one request so dual
|
||||
// consumers (and any accidental parallel fan-out) share one provider call.
|
||||
group singleflight.Group
|
||||
}
|
||||
|
||||
// withWorkflowMetaCache returns a context carrying a mutable per-request cache
|
||||
// for GetWorkflowMeta results. Nested calls reuse the existing cache.
|
||||
func withWorkflowMetaCache(ctx context.Context) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
if _, ok := ctx.Value(workflowMetaCacheKey{}).(*workflowMetaCache); ok {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, workflowMetaCacheKey{}, &workflowMetaCache{
|
||||
byID: make(map[string]workflowMetaCacheEntry),
|
||||
})
|
||||
}
|
||||
|
||||
// getWorkflowMeta loads workflow agent-profile id + prompt via WorkflowStepGetter,
|
||||
// caching the result on ctx when seeded with withWorkflowMetaCache.
|
||||
func (s *Service) getWorkflowMeta(ctx context.Context, workflowID string) (WorkflowMeta, error) {
|
||||
if s.workflowStepGetter == nil || workflowID == "" {
|
||||
return WorkflowMeta{}, nil
|
||||
}
|
||||
|
||||
cache, ok := ctx.Value(workflowMetaCacheKey{}).(*workflowMetaCache)
|
||||
if !ok {
|
||||
return s.workflowStepGetter.GetWorkflowMeta(ctx, workflowID)
|
||||
}
|
||||
|
||||
cache.mu.Lock()
|
||||
if entry, hit := cache.byID[workflowID]; hit {
|
||||
cache.mu.Unlock()
|
||||
return entry.meta, entry.err
|
||||
}
|
||||
cache.mu.Unlock()
|
||||
|
||||
// Coalesce concurrent same-ID loads; first writer still wins on the map.
|
||||
v, _, _ := cache.group.Do(workflowID, func() (any, error) {
|
||||
cache.mu.Lock()
|
||||
if entry, hit := cache.byID[workflowID]; hit {
|
||||
cache.mu.Unlock()
|
||||
return entry, nil
|
||||
}
|
||||
cache.mu.Unlock()
|
||||
|
||||
meta, loadErr := s.workflowStepGetter.GetWorkflowMeta(ctx, workflowID)
|
||||
entry := workflowMetaCacheEntry{meta: meta, err: loadErr}
|
||||
cache.mu.Lock()
|
||||
if existing, hit := cache.byID[workflowID]; hit {
|
||||
cache.mu.Unlock()
|
||||
return existing, nil
|
||||
}
|
||||
cache.byID[workflowID] = entry
|
||||
cache.mu.Unlock()
|
||||
return entry, nil
|
||||
})
|
||||
entry := v.(workflowMetaCacheEntry)
|
||||
return entry.meta, entry.err
|
||||
}
|
||||
@@ -2,8 +2,11 @@ package orchestrator
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
@@ -407,3 +410,283 @@ func TestApplyWorkflowAndPlanMode_PassthroughSkipsReferenceExpansion(t *testing.
|
||||
t.Fatalf("expected expander not to be invoked for a passthrough session, got %d calls", len(expander.calls))
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetWorkflowMeta_CachesPerRequest(t *testing.T) {
|
||||
stepGetter := newMockStepGetter()
|
||||
stepGetter.workflowAgentProfileID = "profile-wf"
|
||||
stepGetter.workflowPrompts["wf-1"] = "Keep CI green."
|
||||
svc := createTestService(setupTestRepo(t), stepGetter, newMockTaskRepo())
|
||||
|
||||
ctx := withWorkflowMetaCache(context.Background())
|
||||
step := &wfmodels.WorkflowStep{
|
||||
ID: "step-1",
|
||||
WorkflowID: "wf-1",
|
||||
// No step override — forces workflow default profile lookup.
|
||||
}
|
||||
|
||||
if got := svc.resolveStepAgentProfile(ctx, step); got != "profile-wf" {
|
||||
t.Fatalf("resolveStepAgentProfile() = %q, want profile-wf", got)
|
||||
}
|
||||
if stepGetter.metaCalls() != 1 {
|
||||
t.Fatalf("after profile resolve: GetWorkflowMeta calls = %d, want 1", stepGetter.metaCalls())
|
||||
}
|
||||
|
||||
got := svc.buildWorkflowPrompt(ctx, "base", step, "task-1", "session-1", false)
|
||||
want := "## Workflow instructions\n\nKeep CI green.\n\n<!-- /workflow-instructions -->\n\nbase"
|
||||
if got != want {
|
||||
t.Fatalf("buildWorkflowPrompt() = %q, want %q", got, want)
|
||||
}
|
||||
if stepGetter.metaCalls() != 1 {
|
||||
t.Fatalf("after prompt build: GetWorkflowMeta calls = %d, want 1 (cached)", stepGetter.metaCalls())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetWorkflowMeta_WithoutCacheHitsProviderTwice(t *testing.T) {
|
||||
stepGetter := newMockStepGetter()
|
||||
stepGetter.workflowAgentProfileID = "profile-wf"
|
||||
stepGetter.workflowPrompts["wf-1"] = "Keep CI green."
|
||||
svc := createTestService(setupTestRepo(t), stepGetter, newMockTaskRepo())
|
||||
|
||||
ctx := context.Background()
|
||||
step := &wfmodels.WorkflowStep{ID: "step-1", WorkflowID: "wf-1"}
|
||||
|
||||
_ = svc.resolveStepAgentProfile(ctx, step)
|
||||
_ = svc.buildWorkflowPrompt(ctx, "base", step, "task-1", "session-1", false)
|
||||
|
||||
if stepGetter.metaCalls() != 2 {
|
||||
t.Fatalf("without cache: GetWorkflowMeta calls = %d, want 2", stepGetter.metaCalls())
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessOnEnter_SharesWorkflowMetaCache(t *testing.T) {
|
||||
// processOnEnter seeds withWorkflowMetaCache; profile switch + prompt build
|
||||
// must share one GetWorkflowMeta read. Cover the pure helpers with the same
|
||||
// ctx seeding pattern processOnEnter uses.
|
||||
stepGetter := newMockStepGetter()
|
||||
stepGetter.workflowAgentProfileID = "profile-wf"
|
||||
stepGetter.workflowPrompts["wf-1"] = "Rule one."
|
||||
svc := createTestService(setupTestRepo(t), stepGetter, newMockTaskRepo())
|
||||
|
||||
ctx := withWorkflowMetaCache(context.Background())
|
||||
step := &wfmodels.WorkflowStep{
|
||||
ID: "step-1",
|
||||
WorkflowID: "wf-1",
|
||||
Prompt: "Do the work.",
|
||||
}
|
||||
|
||||
if got := svc.resolveStepAgentProfile(ctx, step); got != "profile-wf" {
|
||||
t.Fatalf("resolveStepAgentProfile() = %q, want profile-wf", got)
|
||||
}
|
||||
got := svc.buildWorkflowPrompt(ctx, "base", step, "task-1", "session-1", false)
|
||||
if !strings.Contains(got, "Rule one.") {
|
||||
t.Fatalf("expected workflow instructions in prompt, got %q", got)
|
||||
}
|
||||
if stepGetter.metaCalls() != 1 {
|
||||
t.Fatalf("shared cache path: GetWorkflowMeta calls = %d, want 1", stepGetter.metaCalls())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetWorkflowMeta_ErrorIsCachedAndFallsBack(t *testing.T) {
|
||||
// Failure semantics: log + empty profile / omit instructions. Errors are also
|
||||
// cached so dual-consumer step entry does not double-hit a broken provider.
|
||||
stepGetter := newMockStepGetter()
|
||||
stepGetter.workflowMetaErr = errors.New("sqlite: no such workflow")
|
||||
svc := createTestService(setupTestRepo(t), stepGetter, newMockTaskRepo())
|
||||
|
||||
ctx := withWorkflowMetaCache(context.Background())
|
||||
step := &wfmodels.WorkflowStep{ID: "step-1", WorkflowID: "wf-missing"}
|
||||
|
||||
if got := svc.resolveStepAgentProfile(ctx, step); got != "" {
|
||||
t.Fatalf("resolveStepAgentProfile on error = %q, want empty fallback", got)
|
||||
}
|
||||
got := svc.buildWorkflowPrompt(ctx, "base", step, "task-1", "session-1", false)
|
||||
if got != "base" {
|
||||
t.Fatalf("buildWorkflowPrompt on error = %q, want base (instructions omitted)", got)
|
||||
}
|
||||
if stepGetter.metaCalls() != 1 {
|
||||
t.Fatalf("error path GetWorkflowMeta calls = %d, want 1 (cached error)", stepGetter.metaCalls())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetWorkflowMeta_EmptyWorkflowIDSkipsProvider(t *testing.T) {
|
||||
stepGetter := newMockStepGetter()
|
||||
stepGetter.workflowAgentProfileID = "profile-wf"
|
||||
svc := createTestService(setupTestRepo(t), stepGetter, newMockTaskRepo())
|
||||
|
||||
ctx := withWorkflowMetaCache(context.Background())
|
||||
meta, err := svc.getWorkflowMeta(ctx, "")
|
||||
if err != nil {
|
||||
t.Fatalf("getWorkflowMeta(\"\") error = %v", err)
|
||||
}
|
||||
if meta != (WorkflowMeta{}) {
|
||||
t.Fatalf("getWorkflowMeta(\"\") = %+v, want zero value", meta)
|
||||
}
|
||||
if stepGetter.metaCalls() != 0 {
|
||||
t.Fatalf("empty workflowID should not hit provider, calls = %d", stepGetter.metaCalls())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetWorkflowMeta_NilGetterSkipsProvider(t *testing.T) {
|
||||
// Use a bare Service so workflowStepGetter is a true nil interface.
|
||||
// createTestService(nil) would store a typed-nil *mockStepGetter, which is
|
||||
// not == nil and would panic on method call — not a production path.
|
||||
svc := &Service{}
|
||||
ctx := withWorkflowMetaCache(context.Background())
|
||||
|
||||
meta, err := svc.getWorkflowMeta(ctx, "wf-1")
|
||||
if err != nil {
|
||||
t.Fatalf("getWorkflowMeta with nil getter error = %v", err)
|
||||
}
|
||||
if meta != (WorkflowMeta{}) {
|
||||
t.Fatalf("getWorkflowMeta with nil getter = %+v, want zero", meta)
|
||||
}
|
||||
|
||||
step := &wfmodels.WorkflowStep{ID: "step-1", WorkflowID: "wf-1"}
|
||||
if got := svc.resolveStepAgentProfile(ctx, step); got != "" {
|
||||
t.Fatalf("resolveStepAgentProfile with nil getter = %q, want empty", got)
|
||||
}
|
||||
if got := svc.buildWorkflowPrompt(ctx, "base", step, "task-1", "session-1", false); got != "base" {
|
||||
t.Fatalf("buildWorkflowPrompt with nil getter = %q, want base", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetWorkflowMeta_NestedCacheReusesSameMap(t *testing.T) {
|
||||
stepGetter := newMockStepGetter()
|
||||
stepGetter.workflowAgentProfileID = "profile-wf"
|
||||
stepGetter.workflowPrompts["wf-1"] = "Keep CI green."
|
||||
svc := createTestService(setupTestRepo(t), stepGetter, newMockTaskRepo())
|
||||
|
||||
outer := withWorkflowMetaCache(context.Background())
|
||||
inner := withWorkflowMetaCache(outer) // nested seed must be a no-op
|
||||
if outer != inner {
|
||||
t.Fatal("nested withWorkflowMetaCache must reuse the existing cache context")
|
||||
}
|
||||
|
||||
step := &wfmodels.WorkflowStep{ID: "step-1", WorkflowID: "wf-1"}
|
||||
_ = svc.resolveStepAgentProfile(outer, step)
|
||||
_ = svc.buildWorkflowPrompt(inner, "base", step, "task-1", "session-1", false)
|
||||
if stepGetter.metaCalls() != 1 {
|
||||
t.Fatalf("nested cache GetWorkflowMeta calls = %d, want 1", stepGetter.metaCalls())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetWorkflowMeta_CachesPerWorkflowID(t *testing.T) {
|
||||
stepGetter := newMockStepGetter()
|
||||
stepGetter.workflowAgentProfileID = "profile-wf"
|
||||
stepGetter.workflowPrompts["wf-a"] = "Rule A"
|
||||
stepGetter.workflowPrompts["wf-b"] = "Rule B"
|
||||
svc := createTestService(setupTestRepo(t), stepGetter, newMockTaskRepo())
|
||||
|
||||
ctx := withWorkflowMetaCache(context.Background())
|
||||
stepA := &wfmodels.WorkflowStep{ID: "s-a", WorkflowID: "wf-a"}
|
||||
stepB := &wfmodels.WorkflowStep{ID: "s-b", WorkflowID: "wf-b"}
|
||||
|
||||
gotA := svc.buildWorkflowPrompt(ctx, "base", stepA, "task-1", "session-1", false)
|
||||
gotB := svc.buildWorkflowPrompt(ctx, "base", stepB, "task-1", "session-1", false)
|
||||
if !strings.Contains(gotA, "Rule A") || !strings.Contains(gotB, "Rule B") {
|
||||
t.Fatalf("expected distinct prompts per workflow; A=%q B=%q", gotA, gotB)
|
||||
}
|
||||
// Two workflows → two provider hits; second lookup of A must still be cached.
|
||||
_ = svc.buildWorkflowPrompt(ctx, "again", stepA, "task-1", "session-1", false)
|
||||
if stepGetter.metaCalls() != 2 {
|
||||
t.Fatalf("per-id cache calls = %d, want 2", stepGetter.metaCalls())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetWorkflowMeta_WhitespacePromptOmitsInstructions(t *testing.T) {
|
||||
stepGetter := newMockStepGetter()
|
||||
stepGetter.workflowAgentProfileID = "profile-wf"
|
||||
stepGetter.workflowPrompts["wf-1"] = " \n\t "
|
||||
svc := createTestService(setupTestRepo(t), stepGetter, newMockTaskRepo())
|
||||
|
||||
ctx := withWorkflowMetaCache(context.Background())
|
||||
step := &wfmodels.WorkflowStep{ID: "step-1", WorkflowID: "wf-1"}
|
||||
|
||||
// Profile still resolves from the same meta read; only prompt is trimmed away.
|
||||
if got := svc.resolveStepAgentProfile(ctx, step); got != "profile-wf" {
|
||||
t.Fatalf("resolveStepAgentProfile() = %q, want profile-wf", got)
|
||||
}
|
||||
got := svc.buildWorkflowPrompt(ctx, "base", step, "task-1", "session-1", false)
|
||||
if got != "base" {
|
||||
t.Fatalf("whitespace prompt must omit instructions, got %q", got)
|
||||
}
|
||||
if stepGetter.metaCalls() != 1 {
|
||||
t.Fatalf("whitespace path calls = %d, want 1", stepGetter.metaCalls())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetWorkflowMeta_StepProfileOverrideStillReadsPromptOnce(t *testing.T) {
|
||||
// Step agent override short-circuits profile resolution before GetWorkflowMeta,
|
||||
// but prompt build still needs one read. Seeded cache + one call is enough.
|
||||
stepGetter := newMockStepGetter()
|
||||
stepGetter.workflowAgentProfileID = "profile-wf"
|
||||
stepGetter.workflowPrompts["wf-1"] = "Keep CI green."
|
||||
svc := createTestService(setupTestRepo(t), stepGetter, newMockTaskRepo())
|
||||
|
||||
ctx := withWorkflowMetaCache(context.Background())
|
||||
step := &wfmodels.WorkflowStep{
|
||||
ID: "step-1",
|
||||
WorkflowID: "wf-1",
|
||||
AgentProfileID: "profile-step",
|
||||
}
|
||||
|
||||
if got := svc.resolveStepAgentProfile(ctx, step); got != "profile-step" {
|
||||
t.Fatalf("resolveStepAgentProfile() = %q, want step override", got)
|
||||
}
|
||||
if stepGetter.metaCalls() != 0 {
|
||||
t.Fatalf("step override should skip meta for profile, calls = %d", stepGetter.metaCalls())
|
||||
}
|
||||
got := svc.buildWorkflowPrompt(ctx, "base", step, "task-1", "session-1", false)
|
||||
if !strings.Contains(got, "Keep CI green.") {
|
||||
t.Fatalf("expected workflow prompt despite step profile override, got %q", got)
|
||||
}
|
||||
if stepGetter.metaCalls() != 1 {
|
||||
t.Fatalf("prompt-only path calls = %d, want 1", stepGetter.metaCalls())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetWorkflowMeta_ConcurrentReadersSinglePopulate(t *testing.T) {
|
||||
// singleflight + map cache: concurrent same-ID dual consumers share one load.
|
||||
stepGetter := newMockStepGetter()
|
||||
stepGetter.workflowAgentProfileID = "profile-wf"
|
||||
stepGetter.workflowPrompts["wf-1"] = "Keep CI green."
|
||||
// Hold the first provider call long enough that sibling goroutines enter
|
||||
// singleflight.Do while it is still in flight.
|
||||
stepGetter.workflowMetaDelay = 25 * time.Millisecond
|
||||
svc := createTestService(setupTestRepo(t), stepGetter, newMockTaskRepo())
|
||||
|
||||
ctx := withWorkflowMetaCache(context.Background())
|
||||
step := &wfmodels.WorkflowStep{ID: "step-1", WorkflowID: "wf-1"}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
const n = 32
|
||||
wg.Add(n)
|
||||
errs := make(chan string, n)
|
||||
for i := 0; i < n; i++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if got := svc.resolveStepAgentProfile(ctx, step); got != "profile-wf" {
|
||||
errs <- "bad profile: " + got
|
||||
return
|
||||
}
|
||||
prompt := svc.buildWorkflowPrompt(ctx, "base", step, "task-1", "session-1", false)
|
||||
if !strings.Contains(prompt, "Keep CI green.") {
|
||||
errs <- "missing prompt: " + prompt
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for msg := range errs {
|
||||
t.Error(msg)
|
||||
}
|
||||
if calls := stepGetter.metaCalls(); calls != 1 {
|
||||
t.Fatalf("concurrent same-ID load: GetWorkflowMeta calls = %d, want 1", calls)
|
||||
}
|
||||
// Follow-up after quiescence must stay on the map cache.
|
||||
_ = svc.resolveStepAgentProfile(ctx, step)
|
||||
_ = svc.buildWorkflowPrompt(ctx, "base", step, "task-1", "session-1", false)
|
||||
if stepGetter.metaCalls() != 1 {
|
||||
t.Fatalf("post-storm cache miss: calls = %d, want 1", stepGetter.metaCalls())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -169,23 +169,44 @@ func (s *Service) GetNextStepByPosition(ctx context.Context, workflowID string,
|
||||
return nil, nil // No next step found (current step is the last one)
|
||||
}
|
||||
|
||||
// WorkflowMeta is the subset of workflow fields needed at step entry
|
||||
// (agent profile default + optional workflow-level prompt).
|
||||
type WorkflowMeta struct {
|
||||
AgentProfileID string
|
||||
Prompt string
|
||||
}
|
||||
|
||||
// GetWorkflowMeta returns agent profile id and prompt for a workflow in one
|
||||
// provider read. Callers that previously stacked GetWorkflowAgentProfileID and
|
||||
// GetWorkflowPrompt should use this instead.
|
||||
func (s *Service) GetWorkflowMeta(ctx context.Context, workflowID string) (WorkflowMeta, error) {
|
||||
wf, err := s.workflowProvider.GetWorkflow(ctx, workflowID)
|
||||
if err != nil {
|
||||
return WorkflowMeta{}, err
|
||||
}
|
||||
return WorkflowMeta{
|
||||
AgentProfileID: wf.AgentProfileID,
|
||||
Prompt: wf.Prompt,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// GetWorkflowAgentProfileID returns the default agent profile ID for a workflow.
|
||||
func (s *Service) GetWorkflowAgentProfileID(ctx context.Context, workflowID string) (string, error) {
|
||||
wf, err := s.workflowProvider.GetWorkflow(ctx, workflowID)
|
||||
meta, err := s.GetWorkflowMeta(ctx, workflowID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return wf.AgentProfileID, nil
|
||||
return meta.AgentProfileID, nil
|
||||
}
|
||||
|
||||
// GetWorkflowPrompt returns the optional workflow-level agent instructions.
|
||||
// Empty string means the workflow has no prompt configured.
|
||||
func (s *Service) GetWorkflowPrompt(ctx context.Context, workflowID string) (string, error) {
|
||||
wf, err := s.workflowProvider.GetWorkflow(ctx, workflowID)
|
||||
meta, err := s.GetWorkflowMeta(ctx, workflowID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return wf.Prompt, nil
|
||||
return meta.Prompt, nil
|
||||
}
|
||||
|
||||
// GetPreviousStepByPosition returns the previous step before the given position for a workflow.
|
||||
|
||||
@@ -48,7 +48,8 @@ func insertWorkflow(t *testing.T, db *sqlx.DB, id, name string) {
|
||||
|
||||
// mockWorkflowProvider implements WorkflowProvider with in-memory state for tests.
|
||||
type mockWorkflowProvider struct {
|
||||
workflows []*taskmodels.Workflow
|
||||
workflows []*taskmodels.Workflow
|
||||
getWorkflowCalls int
|
||||
}
|
||||
|
||||
func (m *mockWorkflowProvider) ListWorkflows(_ context.Context, workspaceID string, includeHidden bool) ([]*taskmodels.Workflow, error) {
|
||||
@@ -66,6 +67,7 @@ func (m *mockWorkflowProvider) ListWorkflows(_ context.Context, workspaceID stri
|
||||
}
|
||||
|
||||
func (m *mockWorkflowProvider) GetWorkflow(_ context.Context, id string) (*taskmodels.Workflow, error) {
|
||||
m.getWorkflowCalls++
|
||||
for _, wf := range m.workflows {
|
||||
if wf.ID == id {
|
||||
return wf, nil
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
taskmodels "github.com/kandev/kandev/internal/task/models"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetWorkflowMeta_ReturnsProfileAndPromptInOneRead(t *testing.T) {
|
||||
svc, _ := setupTestService(t)
|
||||
provider := &mockWorkflowProvider{}
|
||||
now := time.Now().UTC()
|
||||
provider.workflows = append(provider.workflows, &taskmodels.Workflow{
|
||||
ID: "wf-meta",
|
||||
WorkspaceID: "ws-1",
|
||||
Name: "Meta",
|
||||
AgentProfileID: "profile-42",
|
||||
Prompt: "Always keep CI green.",
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
})
|
||||
svc.SetWorkflowProvider(provider)
|
||||
|
||||
meta, err := svc.GetWorkflowMeta(context.Background(), "wf-meta")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "profile-42", meta.AgentProfileID)
|
||||
assert.Equal(t, "Always keep CI green.", meta.Prompt)
|
||||
assert.Equal(t, 1, provider.getWorkflowCalls)
|
||||
|
||||
// Thin wrappers share the batched path but each call still hits the provider
|
||||
// once when no request cache is present (as expected outside orchestrator).
|
||||
profileID, err := svc.GetWorkflowAgentProfileID(context.Background(), "wf-meta")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "profile-42", profileID)
|
||||
|
||||
prompt, err := svc.GetWorkflowPrompt(context.Background(), "wf-meta")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Always keep CI green.", prompt)
|
||||
assert.Equal(t, 3, provider.getWorkflowCalls)
|
||||
}
|
||||
|
||||
func TestGetWorkflowMeta_PropagatesMissingWorkflow(t *testing.T) {
|
||||
svc, _ := setupTestService(t)
|
||||
provider := &mockWorkflowProvider{}
|
||||
svc.SetWorkflowProvider(provider)
|
||||
|
||||
_, err := svc.GetWorkflowMeta(context.Background(), "missing")
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, 1, provider.getWorkflowCalls)
|
||||
}
|
||||
Reference in New Issue
Block a user