Propagate active call cancellation into component host imports

README.md

@@ -161,13 +161,14 @@ callback first, then close the cache, then close the Wago runtime.
 
 ## Authorities
 
-The provider asks for three required authorities:
+The provider asks for four required authorities:
 
 | Authority | Why it is needed |
 | --- | --- |
 | `core.module.compile` | Compile core Wasm modules embedded in a component. |
 | `core.instance.instantiate` | Instantiate and own the core-module graph, within reviewed positive instance and memory limits. |
 | `core.funcref.create` | Build typed host references for Canonical ABI bridges. |
+| `host.caller.identify` | Observe cancellation and deadlines of the active synchronous host call. |
 
 These handles do not expose plugin registration, runtime policy, hooks, or
 arbitrary runtime lifecycle control. A user may narrow the requested positive

component_test.go

@@ -57,7 +57,7 @@ func consumerProvider(ref **wagoplugin.Ref[component.Service]) wago.PluginProvid
 	}
 }
 
-func pluginSet(t *testing.T, providers []wago.PluginProvider, config json.RawMessage) wago.PluginSet {
+func pluginSet(t testing.TB, providers []wago.PluginProvider, config json.RawMessage) wago.PluginSet {
 	t.Helper()
 	set := wago.PluginSet{Providers: providers}
 	for _, provider := range providers {
@@ -97,7 +97,7 @@ func pluginSet(t *testing.T, providers []wago.PluginProvider, config json.RawMes
 	return set
 }
 
-func loadService(t *testing.T, config json.RawMessage) (*wago.Runtime, *wagoplugin.Ref[component.Service]) {
+func loadService(t testing.TB, config json.RawMessage) (*wago.Runtime, *wagoplugin.Ref[component.Service]) {
 	t.Helper()
 	var ref *wagoplugin.Ref[component.Service]
 	providers := []wago.PluginProvider{component.Provider(), consumerProvider(&ref)}
@@ -293,6 +293,7 @@ func TestDefinitionUsesExactAuthoritiesAndVersionedContract(t *testing.T) {
 		wago.AuthorityCoreModuleCompile,
 		wago.AuthorityCoreInstanceInstantiate,
 		wago.AuthorityCoreFuncRefCreate,
+		wago.AuthorityHostCallerIdentify,
 	}
 	if len(definition.Authorities) != len(wantAuthorities) {
 		t.Fatalf("authorities = %#v", definition.Authorities)

host_call_context_test.go

@@ -0,0 +1,200 @@
+package component_test
+
+import (
+	"context"
+	_ "embed"
+	"errors"
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	component "github.com/wago-org/component-model"
+)
+
+//go:embed testdata/host_call_context.wasm
+var hostCallContextComponent []byte
+
+//go:embed testdata/host_call_context_forwarded.wasm
+var hostCallContextForwardedComponent []byte
+
+type hostContextValueKey struct{}
+
+func TestHostImportReceivesActiveCallCancellation(t *testing.T) {
+	for _, fixture := range []struct {
+		name string
+		wasm []byte
+	}{
+		{"direct", hostCallContextComponent},
+		{"forwarded", hostCallContextForwardedComponent},
+	} {
+		t.Run(fixture.name, func(t *testing.T) {
+			for _, deadline := range []bool{false, true} {
+				name := "cancellation"
+				if deadline {
+					name = "deadline"
+				}
+				t.Run(name, func(t *testing.T) {
+					rt, ref := loadService(t, nil)
+					defer rt.Close()
+					instantiateCtx := context.WithValue(context.Background(), hostContextValueKey{}, "instance")
+					observed := make(chan context.Context, 1)
+					interrupted := make(chan error, 1)
+					rescue := make(chan struct{})
+					var rescueOnce sync.Once
+					release := func() { rescueOnce.Do(func() { close(rescue) }) }
+					defer release()
+					fn := func(ctx context.Context, _ []component.Value) ([]component.Value, error) {
+						observed <- ctx
+						select {
+						case <-ctx.Done():
+							interrupted <- ctx.Err()
+							return nil, ctx.Err()
+						case <-rescue:
+							return nil, nil
+						}
+					}
+					err := ref.With(func(service component.Service) error {
+						return service.WithInstance(instantiateCtx, fixture.wasm, func(in *component.Instance) error {
+							var callCtx context.Context
+							var cancel context.CancelFunc
+							wantErr := error(context.Canceled)
+							if deadline {
+								callCtx, cancel = context.WithTimeout(context.Background(), 100*time.Millisecond)
+								wantErr = context.DeadlineExceeded
+							} else {
+								callCtx, cancel = context.WithCancel(context.Background())
+							}
+							defer cancel()
+							callCtx = context.WithValue(callCtx, hostContextValueKey{}, "call")
+							done := make(chan error, 1)
+							go func() { _, err := in.Call(callCtx, "run"); done <- err }()
+							var hostCtx context.Context
+							select {
+							case hostCtx = <-observed:
+							case err := <-done:
+								return fmt.Errorf("call returned before host import: %v", err)
+							}
+							wantDeadline, wantHasDeadline := callCtx.Deadline()
+							gotDeadline, gotHasDeadline := hostCtx.Deadline()
+							if !deadline {
+								cancel()
+							}
+							var hostErr error
+							select {
+							case hostErr = <-interrupted:
+							case <-time.After(time.Second):
+								release()
+							}
+							callErr := <-done
+							if !errors.Is(hostErr, wantErr) {
+								return fmt.Errorf("active host wait error = %v, want %v (call returned %v)", hostErr, wantErr, callErr)
+							}
+							if !errors.Is(callErr, wantErr) {
+								return fmt.Errorf("component call error = %v, want %v", callErr, wantErr)
+							}
+							if gotHasDeadline != wantHasDeadline || !gotDeadline.Equal(wantDeadline) {
+								return fmt.Errorf("host deadline = %v/%v, want %v/%v", gotDeadline, gotHasDeadline, wantDeadline, wantHasDeadline)
+							}
+							if got := hostCtx.Value(hostContextValueKey{}); got != "instance" {
+								return fmt.Errorf("host context value = %v, want preserved instance value", got)
+							}
+							return nil
+						}, component.WithImport("test:context/host", "wait", fn, nil, nil))
+					})
+					if err != nil {
+						t.Fatal(err)
+					}
+				})
+			}
+		})
+	}
+}
+
+func TestHostImportContextExpiresAfterCallback(t *testing.T) {
+	for _, query := range []bool{false, true} {
+		t.Run(fmt.Sprintf("queried=%v", query), func(t *testing.T) {
+			testHostImportContextExpiresAfterCallback(t, query)
+		})
+	}
+}
+
+func testHostImportContextExpiresAfterCallback(t *testing.T, query bool) {
+	rt, ref := loadService(t, nil)
+	defer rt.Close()
+	var retained context.Context
+	fn := func(ctx context.Context, _ []component.Value) ([]component.Value, error) {
+		retained = ctx
+		if !query {
+			return nil, nil
+		}
+		// Query concurrently to verify the bridge publishes one context safely.
+		var queries sync.WaitGroup
+		for i := 0; i < 8; i++ {
+			queries.Add(1)
+			go func() {
+				defer queries.Done()
+				ctx.Done()
+				ctx.Deadline()
+				ctx.Err()
+			}()
+		}
+		queries.Wait()
+		return nil, nil
+	}
+	err := ref.With(func(service component.Service) error {
+		return service.WithInstance(context.Background(), hostCallContextComponent, func(in *component.Instance) error {
+			values, err := in.Call(context.Background(), "run")
+			if err != nil || len(values) != 1 || values[0] != uint32(7) {
+				return fmt.Errorf("successful host call = %v, %v", values, err)
+			}
+			select {
+			case <-retained.Done():
+			default:
+				return errors.New("retained host context remained live after callback exit")
+			}
+			if err := retained.Err(); !errors.Is(err, context.Canceled) {
+				return fmt.Errorf("retained context error = %v, want cancellation", err)
+			}
+			return nil
+		}, component.WithImport("test:context/host", "wait", fn, nil, nil))
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func BenchmarkHostImportCallContext(b *testing.B) {
+	for _, query := range []bool{false, true} {
+		name := "unused"
+		if query {
+			name = "queried"
+		}
+		b.Run(name, func(b *testing.B) {
+			rt, ref := loadService(b, nil)
+			defer rt.Close()
+			fn := func(ctx context.Context, _ []component.Value) ([]component.Value, error) {
+				if query {
+					ctx.Done()
+				}
+				return nil, nil
+			}
+			err := ref.With(func(service component.Service) error {
+				return service.WithInstance(context.Background(), hostCallContextComponent, func(in *component.Instance) error {
+					b.ReportAllocs()
+					b.ResetTimer()
+					for i := 0; i < b.N; i++ {
+						if _, err := in.Call(context.Background(), "run"); err != nil {
+							return err
+						}
+					}
+					b.StopTimer()
+					return nil
+				}, component.WithImport("test:context/host", "wait", fn, nil, nil))
+			})
+			if err != nil {
+				b.Fatal(err)
+			}
+		})
+	}
+}

internal/engine/engine.go

@@ -10,6 +10,8 @@ import (
 	"fmt"
 	"math"
 	goruntime "runtime"
+	"sync"
+	"time"
 
 	"github.com/wago-org/component-model/internal/engine/expctxkeys"
 	core "github.com/wago-org/wago"
@@ -144,10 +146,56 @@ type runtimeAdapter struct {
 	compiler     *core.CoreModuleCompiler
 	instantiator *core.CoreInstanceInstantiator
 	funcrefs     *core.CoreFuncRefFactory
+	callers      *core.CallerResolver
 }
 
-func Wrap(compiler *core.CoreModuleCompiler, instantiator *core.CoreInstanceInstantiator, funcrefs *core.CoreFuncRefFactory) Runtime {
-	return &runtimeAdapter{compiler: compiler, instantiator: instantiator, funcrefs: funcrefs}
+func Wrap(compiler *core.CoreModuleCompiler, instantiator *core.CoreInstanceInstantiator, funcrefs *core.CoreFuncRefFactory, callers *core.CallerResolver) Runtime {
+	return &runtimeAdapter{compiler: compiler, instantiator: instantiator, funcrefs: funcrefs, callers: callers}
+}
+
+// hostContextSource keeps link-time values separate from the active invocation.
+// It is shared by a module's import bridges, rather than copied on every call.
+type hostContextSource struct {
+	values  context.Context
+	callers *core.CallerResolver
+}
+
+// hostCallContext resolves the scoped invocation only if an import queries
+// cancellation. Imports that only use values do not allocate Wago's invocation
+// state. Once also makes concurrent Context method calls safe.
+type hostCallContext struct {
+	source *hostContextSource
+	caller core.Caller
+	once   sync.Once
+	active context.Context
+}
+
+var expiredHostContext = func() context.Context {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	return ctx
+}()
+
+func (c *hostCallContext) invocation() context.Context {
+	c.once.Do(func() {
+		ctx, err := c.source.callers.InvocationContext(c.caller)
+		if err != nil {
+			// A first query after callback exit must not revive the old context.
+			ctx = expiredHostContext
+		}
+		c.active = ctx
+	})
+	return c.active
+}
+
+func (c *hostCallContext) Deadline() (time.Time, bool) { return c.invocation().Deadline() }
+func (c *hostCallContext) Done() <-chan struct{}       { return c.invocation().Done() }
+func (c *hostCallContext) Err() error                  { return c.invocation().Err() }
+func (c *hostCallContext) Value(key any) any {
+	if _, ok := key.(activeCallerKey); ok {
+		return c.caller
+	}
+	return c.source.values.Value(key)
 }
 
 type compiledModule struct{ mod *core.Module }
@@ -216,7 +264,7 @@ func (r *runtimeAdapter) instantiateModule(ctx context.Context, c CompiledModule
 	if !ok || cm == nil || cm.mod == nil {
 		return nil, fmt.Errorf("component: compiled module belongs to another runtime")
 	}
-	if r == nil || r.instantiator == nil || r.funcrefs == nil {
+	if r == nil || r.instantiator == nil || r.funcrefs == nil || r.callers == nil {
 		return nil, fmt.Errorf("component: incomplete core execution handles")
 	}
 	instantiateOpts := make([]core.InstantiateOption, 0, len(cm.mod.Imports())+1)
@@ -228,6 +276,7 @@ func (r *runtimeAdapter) instantiateModule(ctx context.Context, c CompiledModule
 		}
 	}
 	resolver := resolverFromContext(ctx)
+	var hostContexts *hostContextSource
 	for _, spec := range cm.mod.Imports() {
 		if resolver == nil {
 			continue
@@ -243,6 +292,9 @@ func (r *runtimeAdapter) instantiateModule(ctx context.Context, c CompiledModule
 				continue
 			}
 			resolvedFuncs[coreImportIdentity{module: spec.Module, name: spec.Name}] = fn
+			if hostContexts == nil {
+				hostContexts = &hostContextSource{values: ctx, callers: r.callers}
+			}
 			var host any
 			if hf, ok := fn.(*hostFunction); ok {
 				host = func(caller core.Caller, call core.HostCall) {
@@ -264,14 +316,15 @@ func (r *runtimeAdapter) instantiateModule(ctx context.Context, c CompiledModule
 					}()
 					stack := make([]uint64, max(len(params), len(results)))
 					copy(stack, params)
-					callCtx := context.WithValue(ctx, activeCallerKey{}, caller)
+					callCtx := &hostCallContext{source: hostContexts, caller: caller}
 					hf.fn.Call(callCtx, callerModule{caller: caller}, stack)
 					copy(results, stack)
 				}
 			} else if wf, ok := fn.(*wasmFunction); ok {
 				host = func(caller core.Caller, call core.HostCall) {
 					params, results := call.ParamSlots(), call.ResultSlots()
-					out, callErr := wf.mod.in.InvokeFromHost(ctx, caller, wf.name, params...)
+					callCtx := &hostCallContext{source: hostContexts, caller: caller}
+					out, callErr := wf.mod.in.InvokeFromHost(callCtx, caller, wf.name, params...)
 					if callErr != nil {
 						panic(core.HostTrap{Err: callErr})
 					}

plugin.go

@@ -82,6 +82,11 @@ func Definition() wago.PluginDefinition {
 				Mode:   wago.AuthorityRequired,
 				Reason: "bridge Canonical ABI lifts and lowers through typed host function references",
 			},
+			{
+				Name:   wago.AuthorityHostCallerIdentify,
+				Mode:   wago.AuthorityRequired,
+				Reason: "observe cancellation and deadlines of the active component host call",
+			},
 		},
 		ConfigSchema: append(json.RawMessage(nil), configSchema...),
 		Provides:     []wago.ContractSpec{Contract.Spec()},
@@ -135,7 +140,11 @@ func (*componentPlugin) Register(reg *wago.Registrar) error {
 	if err != nil {
 		return err
 	}
-	service := &runtimeService{engine: engine.Wrap(compiler, instantiator, funcrefs)}
+	callers, err := reg.HostCallers()
+	if err != nil {
+		return err
+	}
+	service := &runtimeService{engine: engine.Wrap(compiler, instantiator, funcrefs, callers)}
 	return wagoplugin.Provide(reg, Contract, Service(service))
 }
 

testdata/host_call_context.wat

@@ -0,0 +1,14 @@
+(component
+  (import "test:context/host" (instance $host
+    (export "wait" (func))))
+  (alias export $host "wait" (func $wait))
+  (core func $wait (canon lower (func $wait)))
+  (core module $guest
+    (import "host" "wait" (func $wait))
+    (func (export "run") (result i32)
+      call $wait
+      i32.const 7))
+  (core instance $host (export "wait" (func $wait)))
+  (core instance $guest (instantiate $guest (with "host" (instance $host))))
+  (func (export "run") (result u32) (canon lift (core func $guest "run")))
+)