fix: pass the account id together with the resolver state (#30)

Author: vahidlazio

wasm/go-host/main.go

@@ -38,7 +38,7 @@ func main() {
 	defer api.Close(ctx)
 
 	// Set resolver state
-	err = api.SetResolverState(resolverState)
+	err = api.SetResolverState(resolverState, "confidence-demo-june")
 	if err != nil {
 		log.Fatalf("Failed to set resolver state: %v", err)
 	}

wasm/go-host/resolver_api.go

@@ -121,17 +121,17 @@ func (r *ResolverApi) Close(ctx context.Context) {
 }
 
 // SetResolverState sets the resolver state in the WASM module
-func (r *ResolverApi) SetResolverState(state []byte) error {
+func (r *ResolverApi) SetResolverState(state []byte, accountId string) error {
 	ctx := context.Background()
 
-	// Create request wrapper
-	request := &messages.Request{
-		Data: state,
+	// Create ResolverStateRequest message
+	resolverStateRequest := &messages.SetResolverStateRequest{
+		State:     state,
+		AccountId: accountId,
 	}
 
 	// Transfer request to WASM memory
-	req, _ := proto.Marshal(request)
-	reqPtr := r.transfer(req)
+	reqPtr := r.transferRequest(resolverStateRequest)
 
 	// Call the WASM function
 	results, err := r.wasmMsgGuestSetResolverState.Call(ctx, uint64(reqPtr))

wasm/java-host/src/main/java/com/spotify/confidence/wasmresolvepoc/Main.java

@@ -2,6 +2,7 @@
 
 import com.dylibso.chicory.wasm.Parser;
 import com.dylibso.chicory.wasm.WasmModule;
+import com.google.protobuf.ByteString;
 import com.google.protobuf.Struct;
 import com.google.protobuf.Value;
 import com.spotify.confidence.flags.resolver.v1.ResolveFlagsRequest;
@@ -11,7 +12,7 @@
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
- 
+import rust_guest.Messages.SetResolverStateRequest;
 
 public class Main {
 
@@ -31,7 +32,9 @@ public static void main(String[] args) throws IOException {
         Path resolveStatePath = Path.of("../resolver_state.pb");
         byte[] resolveState = Files.readAllBytes(resolveStatePath);
 
-        main.resolverApi.setResolverState(resolveState);
+        main.resolverApi
+                .setResolverState(SetResolverStateRequest.newBuilder()
+                        .setState(ByteString.copyFrom(resolveState)).setAccountId("confidence-demo-june").build());
 
         // Verify RESOLVE_REASON_MATCH reason and non-empty variant for tutorial_visitor
         final ResolveFlagsResponse verifyResp = main.resolverApi.resolve(

wasm/java-host/src/main/java/com/spotify/confidence/wasmresolvepoc/ResolverApi.java

@@ -17,6 +17,7 @@
 import com.spotify.confidence.flags.resolver.v1.ResolveFlagsResponse;
 import com.spotify.confidence.flags.resolver.v1.ResolvedFlag;
 
+import rust_guest.Messages.SetResolverStateRequest;
 import rust_guest.Messages;
 import rust_guest.Messages.ResolveSimpleRequest;
 import rust_guest.Types;
@@ -69,9 +70,9 @@ private Timestamp currentTime(Messages.Void unused) {
     return Timestamp.getDefaultInstance();
   }
 
-  public void setResolverState(byte[] state) {
+  public void setResolverState(SetResolverStateRequest state) {
     final byte[] request = Messages.Request.newBuilder()
-            .setData(ByteString.copyFrom(state))
+            .setData(state.toByteString())
             .build().toByteArray();
     int addr = transfer(request);
     int respPtr = (int) wasmMsgGuestSetResolverState.apply(addr)[0];

wasm/node-host/src/index.ts

@@ -1,7 +1,12 @@
 import fs from 'node:fs';
 import path from 'node:path';
 import { Void } from './proto/messages.ts';
-import { ResolveFlagsRequest, ResolveFlagsResponse, ResolveReason } from './proto/resolver/api.ts';
+import {
+  ResolveFlagsRequest,
+  ResolveFlagsResponse,
+  ResolveReason
+} from './proto/resolver/api.ts';
+import { SetResolverStateRequest } from './proto/messages.js';
 import { Timestamp } from './proto/google/protobuf/timestamp.ts';
 import { ApiBuilder } from './wasm-msg.ts';
 
@@ -16,7 +21,7 @@ const resolverState = fs.readFileSync(path.join(dirname, '../../resolver_state.p
 const wasmModule = new WebAssembly.Module(wasmBuffer);
 
 const api = new ApiBuilder()
-  .guestRaw('set_resolver_state')
+  .guest('set_resolver_state', SetResolverStateRequest, Void, false)
   .guest('resolve', ResolveFlagsRequest, ResolveFlagsResponse, false)
   .host('log_resolve', Void, Void, false, () => ({}))
   .host('log_assign', Void, Void, false, () => ({}))
@@ -26,7 +31,10 @@ const api = new ApiBuilder()
   })
   .build(wasmModule);
 
-api.set_resolver_state(resolverState);
+api.set_resolver_state({
+  state: resolverState,
+  accountId: 'confidence-demo-june'
+});
 
 // Verify MATCH reason and non-empty variant for tutorial_visitor
 {

wasm/proto/messages.proto

@@ -8,6 +8,11 @@ import "google/protobuf/struct.proto";
 
 message Void {}
 
+message SetResolverStateRequest {
+    bytes state = 1;
+    string account_id = 2;
+}
+
 message ResolveSimpleRequest {
     string client_secret = 1;
     google.protobuf.Struct evaluation_context = 2;

wasm/proto/resolver/api.proto

@@ -79,7 +79,6 @@ message ResolvedFlag {
   ResolveReason reason = 5;
 }
 
-
 enum ResolveReason {
   // Unspecified enum.
   RESOLVE_REASON_UNSPECIFIED = 0;

wasm/python-host/main.py

@@ -29,7 +29,7 @@ def main():
 
     # Set resolver state
     try:
-        api.set_resolver_state(resolver_state)
+        api.set_resolver_state(resolver_state, "confidence-demo-june")
     except Exception as e:
         print(f"Failed to set resolver state: {e}")
         sys.exit(1)

wasm/python-host/resolver_api.py

@@ -93,14 +93,14 @@ def current_thread_id() -> int:
         # Instantiate the module with imports
         self.instance = linker.instantiate(self.store, self.module)
 
-    def set_resolver_state(self, state: bytes) -> None:
+    def set_resolver_state(self, state: bytes, account_id: str) -> None:
         """Set the resolver state in the WASM module"""
-        # Create request wrapper
-        request = messages_pb2.Request()
-        request.data = state
-        # Transfer request to WASM memory - ensure it's a bytestring
-        serialized_data: bytes = request.SerializeToString()
-        req_ptr = self._transfer(serialized_data)
+        # Create SetResolverStateRequest proto
+        set_resolver_state_request = messages_pb2.SetResolverStateRequest()
+        set_resolver_state_request.state = state
+        set_resolver_state_request.account_id = account_id
+        # Transfer request to WASM memory
+        req_ptr = self._transfer_request(set_resolver_state_request)
         # Call the WASM function
         results = self.wasm_msg_guest_set_resolver_state(self.store, req_ptr)
         resp_ptr = results

wasm/rust-guest/src/lib.rs

@@ -1,10 +1,10 @@
 use std::cell::RefCell;
 use std::sync::Arc;
 use std::sync::LazyLock;
-use std::sync::OnceLock;
 
 use arc_swap::ArcSwapOption;
 use bytes::Bytes;
+use prost::Message;
 
 #[global_allocator]
 static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;
@@ -24,20 +24,18 @@ use wasm_msg::WasmResult;
 pub mod proto {
     include!(concat!(env!("OUT_DIR"), "/rust_guest.rs"));
 }
-use proto::{ResolveSimpleRequest, Void};
-
-use crate::proto::{LogAssignRequest, LogResolveRequest, SdkId};
+use crate::proto::SetResolverStateRequest;
 use confidence_resolver::{
     proto::{
         confidence::flags::admin::v1::ResolverState as ResolverStatePb,
-        confidence::flags::resolver::v1::resolve_token_v1::AssignedFlag,
         confidence::flags::resolver::v1::{
             ResolveFlagsRequest, ResolveFlagsResponse, ResolvedFlag, Sdk,
         },
         google::{Struct, Timestamp},
     },
     Client, FlagToApply, Host, ResolveReason, ResolvedValue, ResolverState,
 };
+use proto::{ResolveSimpleRequest, Void};
 
 impl Into<proto::FallthroughAssignment>
     for confidence_resolver::proto::confidence::flags::resolver::v1::events::FallthroughAssignment
@@ -53,8 +51,6 @@ impl Into<proto::FallthroughAssignment>
 }
 
 const VOID: Void = Void {};
-
-const ACCOUNT_ID: &str = "confidence-test";
 const ENCRYPTION_KEY: Bytes = Bytes::from_static(&[0; 16]);
 
 // TODO simplify by assuming single threaded?
@@ -180,8 +176,10 @@ fn get_resolver_state() -> Result<Arc<ResolverState>, String> {
 }
 
 wasm_msg_guest! {
-    fn set_resolver_state(request: ResolverStatePb) -> WasmResult<Void> {
-        let new_state = ResolverState::from_proto(request, ACCOUNT_ID)?;
+    fn set_resolver_state(request: SetResolverStateRequest) -> WasmResult<Void> {
+        let state_pb = ResolverStatePb::decode(request.state.as_slice())
+            .map_err(|e| format!("Failed to decode resolver state: {}", e))?;
+        let new_state = ResolverState::from_proto(state_pb, request.account_id.as_str())?;
         RESOLVER_STATE.store(Some(Arc::new(new_state)));
         Ok(VOID)
     }