This commit is contained in:
Quanzheng Long
2026-03-18 11:41:05 -07:00
parent 12ea4c24d5
commit c0bd2dbc07
10 changed files with 505 additions and 61 deletions
+77 -3
View File
@@ -98,13 +98,13 @@ type Context struct {
resumeEvent *WaitEvent
}
func (c *Context) WaitFor(cond AnyOfCondition) (WaitEvent, error) {
func (c *Context) WaitFor(cond AnyOfCondition) (WaitForResult, error) {
if c.resumeEvent != nil {
event := *c.resumeEvent
c.resumeEvent = nil
return event, nil
return waitForResultFromRaw(cond, event), nil
}
return WaitEvent{}, ErrWaitRequested{Condition: cond}
return WaitForResult{}, ErrWaitRequested{Condition: cond}
}
func (c *Context) PublishToChannel(channel string, value any) error {
@@ -390,3 +390,77 @@ func unwrapResumeInput(input any) (any, *WaitEvent) {
}
return rawArg, &event
}
func waitForResultFromRaw(cond AnyOfCondition, event WaitEvent) WaitForResult {
result := WaitForResult{
Conditions: make([]ConditionResult, len(cond.Conditions)),
}
if event.Condition == "timer" {
for i, raw := range cond.Conditions {
if kind, _ := raw["kind"].(string); kind == "timer" {
result.Conditions[i] = ConditionResult{Met: true}
break
}
}
return result
}
if event.Condition != "channel" {
return result
}
if event.Channel == "__any_of__" {
var matched []struct {
Channel string `json:"channel"`
Value any `json:"value"`
}
_ = json.Unmarshal(event.Value, &matched)
cursor := 0
for i, raw := range cond.Conditions {
kind, _ := raw["kind"].(string)
if kind != "channel" || cursor >= len(matched) {
continue
}
channelName, _ := raw["channel"].(string)
if channelName == matched[cursor].Channel {
result.Conditions[i] = ConditionResult{
Met: true,
ChannelName: channelName,
Values: toValues(matched[cursor].Value),
}
cursor++
}
}
return result
}
var value any
_ = json.Unmarshal(event.Value, &value)
for i, raw := range cond.Conditions {
kind, _ := raw["kind"].(string)
if kind != "channel" {
continue
}
channelName, _ := raw["channel"].(string)
if channelName == event.Channel {
result.Conditions[i] = ConditionResult{
Met: true,
ChannelName: channelName,
Values: toValues(value),
}
break
}
}
return result
}
func toValues(value any) []any {
if value == nil {
return []any{}
}
if vals, ok := value.([]any); ok {
return vals
}
return []any{value}
}
+10
View File
@@ -58,6 +58,16 @@ type WaitEvent struct {
Seconds float64 `json:"seconds,omitempty"`
}
type ConditionResult struct {
Met bool `json:"met"`
ChannelName string `json:"channel_name,omitempty"`
Values []any `json:"values,omitempty"`
}
type WaitForResult struct {
Conditions []ConditionResult `json:"conditions"`
}
type Send struct {
Node any
NodeInput any
+12 -3
View File
@@ -37,14 +37,23 @@ func (c *Context) Interrupt(name string) (any, error) {
}
return nil, err
}
if len(event.Value) == 0 {
if len(event.Conditions) == 0 || !event.Conditions[0].Met {
return nil, nil
}
values := event.Conditions[0].Values
if len(values) == 0 {
return nil, nil
}
rawValue := values[0]
valueBytes, err := json.Marshal(rawValue)
if err != nil {
return nil, fmt.Errorf("encode interrupt `%s` value: %w", name, err)
}
var payload interruptPayload
if err := json.Unmarshal(event.Value, &payload); err != nil {
if err := json.Unmarshal(valueBytes, &payload); err != nil {
var value any
if err := json.Unmarshal(event.Value, &value); err != nil {
if err := json.Unmarshal(valueBytes, &value); err != nil {
return nil, fmt.Errorf("decode interrupt `%s` value: %w", name, err)
}
return value, nil
+18 -9
View File
@@ -75,16 +75,25 @@ func (w *lunchWorkflow) waitNode(ctx *ag.Context, _ any, state lunchState) (ag.C
}
output := append([]string(nil), state.Output...)
if event.Condition == "channel" {
payload := ag.DecodeString(event.Value)
switch event.Channel {
case "tool_completion_channel":
output = append(output, "tool: "+payload)
case "subagent_completion_channel":
output = append(output, "sub_agent: "+payload)
case "user_input_channel":
output = append(output, "user_input: "+payload)
hadChannel := false
for _, cond := range event.Conditions {
if !cond.Met || cond.ChannelName == "" {
continue
}
for _, raw := range cond.Values {
payload, _ := raw.(string)
switch cond.ChannelName {
case "tool_completion_channel":
output = append(output, "tool: "+payload)
case "subagent_completion_channel":
output = append(output, "sub_agent: "+payload)
case "user_input_channel":
output = append(output, "user_input: "+payload)
}
}
hadChannel = true
}
if hadChannel {
state.Output = output
return ag.Command{Goto: []ag.Send{{Node: w.llmNode}}, Update: state}, nil
}
+105 -5
View File
@@ -1,7 +1,6 @@
package tests
import (
"encoding/json"
"fmt"
"testing"
@@ -143,7 +142,7 @@ func (w *primitiveWorkflow) startWaitBatchNode(ctx *ag.Context, _ any, state pri
}
func (w *primitiveWorkflow) waitBatchNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
event, err := ctx.WaitFor(ag.AnyOf(ag.ChannelCondition{
result, err := ctx.WaitFor(ag.AnyOf(ag.ChannelCondition{
Channel: "events",
Min: 2,
Max: 4,
@@ -151,9 +150,13 @@ func (w *primitiveWorkflow) waitBatchNode(ctx *ag.Context, _ any, state primitiv
if err != nil {
return ag.Command{}, err
}
var values []string
if err := json.Unmarshal(event.Value, &values); err != nil {
return ag.Command{}, err
if len(result.Conditions) != 1 || !result.Conditions[0].Met {
return ag.Command{}, fmt.Errorf("expected one met condition")
}
values := make([]string, 0, len(result.Conditions[0].Values))
for _, v := range result.Conditions[0].Values {
s, _ := v.(string)
values = append(values, s)
}
state.Count = len(values)
state.Logs = values
@@ -191,3 +194,100 @@ func TestChannelWaitRespectsMaxM(t *testing.T) {
t.Fatalf("unexpected done: %v", result.Done)
}
}
func (w *primitiveWorkflow) startAnyOfTwoChannelsNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
if err := ctx.PublishToChannel("alpha", "a1"); err != nil {
return ag.Command{}, err
}
if err := ctx.PublishToChannel("beta", "b1"); err != nil {
return ag.Command{}, err
}
return ag.Command{
Update: state,
Goto: []ag.Send{
{Node: w.waitAnyOfTwoChannelsNode, NodeInput: nil},
},
}, nil
}
func (w *primitiveWorkflow) waitAnyOfTwoChannelsNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
first, err := ctx.WaitFor(ag.AnyOf(
ag.ChannelCondition{Channel: "alpha"},
ag.ChannelCondition{Channel: "beta"},
))
if err != nil {
return ag.Command{}, err
}
if len(first.Conditions) != 2 {
return ag.Command{}, fmt.Errorf("expected 2 condition results, got %d", len(first.Conditions))
}
if !first.Conditions[0].Met || first.Conditions[0].ChannelName != "alpha" || len(first.Conditions[0].Values) != 1 || first.Conditions[0].Values[0] != "a1" {
return ag.Command{}, fmt.Errorf("unexpected first condition result: %#v", first.Conditions[0])
}
if !first.Conditions[1].Met || first.Conditions[1].ChannelName != "beta" || len(first.Conditions[1].Values) != 1 || first.Conditions[1].Values[0] != "b1" {
return ag.Command{}, fmt.Errorf("unexpected second condition result: %#v", first.Conditions[1])
}
if err := ctx.PublishToChannel("beta", "b2"); err != nil {
return ag.Command{}, err
}
state.Count = 1
state.Logs = []string{
"matched=2",
}
return ag.Command{
Update: state,
Goto: []ag.Send{
{Node: w.verifyBetaAfterAnyOfNode, NodeInput: nil},
},
}, nil
}
func (w *primitiveWorkflow) verifyBetaAfterAnyOfNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
second, err := ctx.WaitFor(ag.AnyOf(
ag.ChannelCondition{Channel: "beta"},
))
if err != nil {
return ag.Command{}, err
}
if len(second.Conditions) != 1 || !second.Conditions[0].Met || second.Conditions[0].ChannelName != "beta" || len(second.Conditions[0].Values) != 1 {
return ag.Command{}, fmt.Errorf("unexpected beta condition result: %#v", second.Conditions)
}
payload, _ := second.Conditions[0].Values[0].(string)
state.Count = 2
state.Logs = append(state.Logs, fmt.Sprintf("beta=%s", payload))
state.Done = "ok"
return ag.Command{Update: state}, nil
}
func TestAnyOfConsumesAllReadyChannels(t *testing.T) {
workflow := &primitiveWorkflow{}
graph := ag.NewAdvancedStateGraph[primitiveState]()
graph.AddAsyncChannel("alpha")
graph.AddAsyncChannel("beta")
graph.AddEntryNode(workflow.startAnyOfTwoChannelsNode)
graph.AddNode(workflow.waitAnyOfTwoChannelsNode)
graph.AddFinishNode(workflow.verifyBetaAfterAnyOfNode)
handler, err := graph.Compile().Start(nil, primitiveState{
Count: 0,
Logs: []string{},
Done: "",
})
if err != nil {
t.Fatalf("start failed: %v", err)
}
result, err := handler.WaitForResult()
if err != nil {
t.Fatalf("result failed: %v", err)
}
if result.Count != 2 {
t.Fatalf("unexpected count: %v", result.Count)
}
if len(result.Logs) != 2 || result.Logs[0] != "matched=2" || result.Logs[1] != "beta=b2" {
t.Fatalf("unexpected logs: %#v", result.Logs)
}
if result.Done != "ok" {
t.Fatalf("unexpected done: %v", result.Done)
}
}